add mcp bridge (#107)

This commit is contained in:
Zhanghaibin
2022-12-09 15:18:49 +08:00
committed by GitHub
parent 4776d62515
commit b08b00f9d5
26 changed files with 4290 additions and 0 deletions

View File

@@ -0,0 +1,171 @@
package address
import (
"fmt"
"io/ioutil"
"math/rand"
"net"
"net/http"
"strings"
"sync"
"time"
"go.uber.org/atomic"
"istio.io/pkg/log"
)
const (
NACOS_PATH = "/nacos/serverlist"
MODULE_HEADER_KEY = "Request-Module"
MODULE_HEADER_VALUE = "Naming"
DEFAULT_INTERVAL = 30 * time.Second
)
type NacosAddressProvider struct {
serverAddr string
nacosAddr string
nacosBackupAddr []string
namespace string
stop chan struct{}
trigger chan struct{}
cond *sync.Cond
isStop *atomic.Bool
mutex *sync.Mutex
}
func NewNacosAddressProvider(serverAddr, namespace string) *NacosAddressProvider {
provider := &NacosAddressProvider{
serverAddr: serverAddr,
namespace: namespace,
stop: make(chan struct{}),
trigger: make(chan struct{}, 1),
cond: sync.NewCond(new(sync.Mutex)),
isStop: atomic.NewBool(false),
mutex: &sync.Mutex{},
}
go provider.Run()
return provider
}
func (p *NacosAddressProvider) Run() {
ticker := time.NewTicker(DEFAULT_INTERVAL)
defer ticker.Stop()
p.addressDiscovery()
for {
select {
case <-p.trigger:
p.addressDiscovery()
case <-ticker.C:
p.addressDiscovery()
case <-p.stop:
return
}
}
}
func (p *NacosAddressProvider) Update(serverAddr, namespace string) {
p.mutex.Lock()
p.serverAddr = serverAddr
p.namespace = namespace
p.mutex.Unlock()
p.addressDiscovery()
}
func (p *NacosAddressProvider) Trigger() {
p.cond.L.Lock()
oldAddr := p.nacosAddr
if len(p.nacosBackupAddr) > 0 {
p.nacosAddr = p.nacosBackupAddr[rand.Intn(len(p.nacosBackupAddr))]
for i := len(p.nacosBackupAddr) - 1; i >= 0; i-- {
if p.nacosBackupAddr[i] == p.nacosAddr {
p.nacosBackupAddr = append(p.nacosBackupAddr[:i], p.nacosBackupAddr[i+1:]...)
}
}
p.nacosBackupAddr = append(p.nacosBackupAddr, oldAddr)
}
p.cond.Broadcast()
p.cond.L.Unlock()
select {
case p.trigger <- struct{}{}:
default:
}
}
func (p *NacosAddressProvider) Stop() {
p.isStop.Store(true)
p.stop <- struct{}{}
}
func (p *NacosAddressProvider) GetNacosAddress(oldAddress string) <-chan string {
addressChan := make(chan string)
go func() {
var addr string
p.cond.L.Lock()
log.Debugf("get nacos address, p.nacosAddr, oldAddress", p.nacosAddr, oldAddress)
for p.nacosAddr == oldAddress || p.nacosAddr == "" {
if p.isStop.Load() {
return
}
p.cond.Wait()
}
addr = p.nacosAddr
p.cond.L.Unlock()
addressChan <- addr
}()
return addressChan
}
func (p *NacosAddressProvider) addressDiscovery() {
p.mutex.Lock()
url := fmt.Sprintf("%s%s?namespace=%s", p.serverAddr, NACOS_PATH, p.namespace)
p.mutex.Unlock()
if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
url = "http://" + url
}
req, err := http.NewRequest("GET", url, nil)
if err != nil {
log.Errorf("create request failed, err:%v, url:%s", err, url)
return
}
req.Header.Add(MODULE_HEADER_KEY, MODULE_HEADER_VALUE)
resp, err := http.DefaultClient.Do(req)
if err != nil {
log.Errorf("get nacos address failed, err:%v, url:%s", err, url)
return
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
log.Errorf("get nacos address failed, statusCode:%d", resp.StatusCode)
return
}
body, _ := ioutil.ReadAll(resp.Body)
addresses := string(body)
addrVec := strings.Fields(addresses)
if len(addrVec) == 0 {
return
}
needUpdate := true
p.cond.L.Lock()
for _, address := range addrVec {
ip := net.ParseIP(address)
if ip == nil {
log.Errorf("ip parse failed, ip:%s", address)
return
}
if p.nacosAddr == address {
needUpdate = false
}
}
p.nacosBackupAddr = addrVec
if needUpdate {
p.nacosAddr = addrVec[rand.Intn(len(addrVec))]
p.cond.Broadcast()
log.Infof("nacos address updated, address:%s", p.nacosAddr)
}
for i := len(p.nacosBackupAddr) - 1; i >= 0; i-- {
if p.nacosBackupAddr[i] == p.nacosAddr {
p.nacosBackupAddr = append(p.nacosBackupAddr[:i], p.nacosBackupAddr[i+1:]...)
}
}
p.cond.L.Unlock()
}

View File

@@ -0,0 +1,287 @@
package address
import (
"net/http"
"net/http/httptest"
"testing"
"time"
)
func setUpServer(status int, body []byte) (string, func()) {
server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(status)
rw.Write(body)
}))
return server.URL, func() {
server.Close()
}
}
func setUpServerWithBodyPtr(status int, body *[]byte) (string, func()) {
server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
rw.WriteHeader(status)
rw.Write(*body)
}))
return server.URL, func() {
server.Close()
}
}
func TestGetNacosAddress(t *testing.T) {
goodURL, goodTearDown := setUpServer(200, []byte("1.1.1.1\n 2.2.2.2"))
defer goodTearDown()
badURL, badTearDown := setUpServer(200, []byte("abc\n 2.2.2.2"))
defer badTearDown()
errURL, errTearDown := setUpServer(503, []byte("1.1.1.1\n 2.2.2.2"))
defer errTearDown()
tests := []struct {
name string
serverAddr string
want []string
}{
{
"good",
goodURL,
[]string{"1.1.1.1", "2.2.2.2"},
},
{
"bad",
badURL,
[]string{},
},
{
"err",
errURL,
[]string{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
provider := NewNacosAddressProvider(tt.serverAddr, "")
timeout := time.NewTicker(1 * time.Second)
var got string
if len(tt.want) == 0 {
select {
case got = <-provider.GetNacosAddress(""):
t.Errorf("GetNacosAddress() = %v, want empty", got)
case <-timeout.C:
return
}
}
select {
case got = <-provider.GetNacosAddress(""):
case <-timeout.C:
t.Error("GetNacosAddress timeout")
}
for _, value := range tt.want {
if got == value {
return
}
}
t.Errorf("GetNacosAddress() = %v, want %v", got, tt.want)
})
}
}
func TestTrigger(t *testing.T) {
body := []byte("1.1.1.1 ")
url, tearDown := setUpServerWithBodyPtr(200, &body)
defer tearDown()
provider := NewNacosAddressProvider(url, "xxxx")
address := <-provider.GetNacosAddress("")
if address != "1.1.1.1" {
t.Errorf("got %s, want %s", address, "1.1.1.1")
}
body = []byte(" 2.2.2.2 ")
tests := []struct {
name string
trigger bool
want string
}{
{
"no trigger",
false,
"1.1.1.1",
},
{
"trigger",
true,
"2.2.2.2",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if tt.trigger {
provider.Trigger()
}
timeout := time.NewTicker(1 * time.Second)
select {
case <-provider.GetNacosAddress("1.1.1.1"):
case <-timeout.C:
}
if provider.nacosAddr != tt.want {
t.Errorf("got %s, want %s", provider.nacosAddr, tt.want)
}
})
}
}
func TestBackup(t *testing.T) {
body := []byte("1.1.1.1 ")
url, tearDown := setUpServerWithBodyPtr(200, &body)
defer tearDown()
provider := NewNacosAddressProvider(url, "xxxx")
address := <-provider.GetNacosAddress("")
if address != "1.1.1.1" {
t.Errorf("got %s, want %s", address, "1.1.1.1")
}
tests := []struct {
name string
oldaddr string
newaddr string
triggerNum int
want string
}{
{
"case1",
"1.1.1.1",
"1.1.1.1\n2.2.2.2",
1,
"2.2.2.2",
},
{
"case2",
"1.1.1.1",
"3.3.3.3 1.1.1.1",
1,
"3.3.3.3",
},
{
"case3",
"1.1.1.1",
"3.3.3.3 1.1.1.1",
2,
"1.1.1.1",
},
{
"case4",
"1.1.1.1",
"3.3.3.3\n 1.1.1.1",
3,
"3.3.3.3",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
provider.nacosAddr = tt.oldaddr
body = []byte(tt.newaddr)
provider.addressDiscovery()
for i := 0; i < tt.triggerNum; i++ {
provider.Trigger()
}
timeout := time.NewTicker(1 * time.Second)
var newAddr string
select {
case newAddr = <-provider.GetNacosAddress(""):
case <-timeout.C:
}
if newAddr != tt.want {
t.Errorf("got %s, want %s", newAddr, tt.want)
}
})
}
}
func TestKeepIp(t *testing.T) {
body := []byte("1.1.1.1")
url, tearDown := setUpServerWithBodyPtr(200, &body)
defer tearDown()
provider := NewNacosAddressProvider(url, "xxxx")
address := <-provider.GetNacosAddress("")
if address != "1.1.1.1" {
t.Errorf("got %s, want %s", address, "1.1.1.1")
}
tests := []struct {
name string
newAddr []byte
want string
}{
{
"add ip",
[]byte("1.1.1.1\n 2.2.2.2"),
"1.1.1.1",
},
{
"remove ip",
[]byte("2.2.2.2"),
"2.2.2.2",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body = tt.newAddr
provider.addressDiscovery()
timeout := time.NewTicker(1 * time.Second)
select {
case <-provider.GetNacosAddress("1.1.1.1"):
case <-timeout.C:
}
if provider.nacosAddr != tt.want {
t.Errorf("got %s, want %s", provider.nacosAddr, tt.want)
}
})
}
}
func TestMultiClient(t *testing.T) {
body := []byte("1.1.1.1")
url, tearDown := setUpServerWithBodyPtr(200, &body)
defer tearDown()
provider := NewNacosAddressProvider(url, "xxxx")
address := <-provider.GetNacosAddress("")
if address != "1.1.1.1" {
t.Errorf("got %s, want %s", address, "1.1.1.1")
}
body = []byte("2.2.2.2")
tests := []struct {
name string
oldAddrs []string
want []string
}{
{
"case1",
[]string{"1.1.1.1", "1.1.1.1"},
[]string{"2.2.2.2", "2.2.2.2"},
},
{
"case2",
[]string{"2.2.2.2", "1.1.1.1"},
[]string{"", "2.2.2.2"},
},
{
"case3",
[]string{"1.1.1.1", "2.2.2.2"},
[]string{"2.2.2.2", ""},
},
{
"case4",
[]string{"2.2.2.2", "2.2.2.2"},
[]string{"", ""},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
provider.addressDiscovery()
for i := 0; i < len(tt.oldAddrs); i++ {
timeout := time.NewTicker(1 * time.Second)
var newaddr string
select {
case newaddr = <-provider.GetNacosAddress(tt.oldAddrs[i]):
case <-timeout.C:
}
if newaddr != tt.want[i] {
t.Errorf("got %s, want %s", newaddr, tt.want[i])
}
}
})
}
}