feat: support Consul as ServiceRegistry

controller/registries/consul/config.go

@@ -45,10 +45,6 @@
 	})
 }
 
-const (
-	defaultToken = ""
-)
-
 type Consul struct {
 	consul.RegistryType
 	logger log.RegistryLogger
@@ -70,11 +66,7 @@
 
 	DataCenter string
 	NameSpace  string
-}
-
-type consulService struct {
-	DataCenter  string
-	ServiceName string
+	Token      string
 }
 
 func (reg *Consul) NewClient(config *consul.Config) (*Client, error) {
@@ -85,7 +77,7 @@
 	clientConfig := consulapi.DefaultConfig()
 	clientConfig.Address = uri.Host
 	clientConfig.Scheme = uri.Scheme
-	clientConfig.Token = defaultToken
+	clientConfig.Token = config.Token
 	clientConfig.Datacenter = config.DataCenter
 
 	client, err := consulapi.NewClient(clientConfig)
@@ -97,9 +89,16 @@
 		consulClient:  client,
 		consulCatalog: client.Catalog(),
 		DataCenter:    config.DataCenter,
+		NameSpace:     config.Namespace,
+		Token:         config.Token,
 	}, nil
 }
 
+type consulService struct {
+	Tag         string
+	ServiceName string
+}
+
 func (reg *Consul) Start(c registrytype.RegistryConfig) error {
 	config := c.(*consul.Config)
 
@@ -108,11 +107,13 @@
 		return err
 	}
 
-	services, err := reg.fetchAllServices(client)
+	reg.client = client
+
+	services, err := reg.fetchAllServices(reg.client)
+
 	if err != nil {
-		return fmt.Errorf("fetch all services error: %v", err)
+		return err
 	}
-	reg.client = client
 
 	reg.watchingServices = services
 
@@ -122,22 +123,29 @@
 	}
 	go func() {
 		reg.logger.Infof("start refreshing services")
-		ticker := time.NewTicker(dur)
-		//q := consulapi.QueryOptions{
-		//    WaitTime: dur,
-		//}
-		defer ticker.Stop()
+		q := &consulapi.QueryOptions{
+			WaitTime:   dur,
+			Namespace:  config.Namespace,
+			Datacenter: config.DataCenter,
+			Token:      config.Token,
+		}
 		for {
+
 			select {
-			case <-ticker.C:
-				err := reg.refresh()
-				if err != nil {
-					reg.logger.Errorf("failed to refresh services, err: %v", err)
-				}
 			case <-reg.done:
 				reg.logger.Infof("stop refreshing services")
 				return
+			default:
+			}
+			services, meta, err := reg.client.consulCatalog.Services(q)
+			if err != nil {
+				reg.logger.Errorf("failed to get services, err: %v", err)
+				time.Sleep(dur)
+				continue
 			}
+			reg.refresh(services)
+
+			q.WaitIndex = meta.LastIndex
 		}
 	}()
 
@@ -160,13 +168,28 @@
 	return nil
 }
 
-func (reg *Consul) refresh() error {
-	return nil
-}
-
 func (reg *Consul) fetchAllServices(client *Client) (map[consulService]bool, error) {
-	fmt.Println(client)
-	return nil, nil
+	q := &consulapi.QueryOptions{}
+	q.Datacenter = client.DataCenter
+	q.Namespace = client.NameSpace
+	q.Token = client.Token
+	services, _, err := client.consulCatalog.Services(q)
+
+	if err != nil {
+		reg.logger.Errorf("failed to get service, err: %v", err)
+		return nil, err
+	}
+	serviceMap := make(map[consulService]bool)
+	for serviceName, tags := range services {
+		for _, tag := range tags {
+			service := consulService{
+				Tag:         tag,
+				ServiceName: serviceName,
+			}
+			serviceMap[service] = true
+		}
+	}
+	return serviceMap, nil
 }
 
 func (reg *Consul) subscribe(serviceName string) error {
@@ -178,3 +201,37 @@
 	fmt.Println(serviceName)
 	return nil
 }
+
+func (reg *Consul) refresh(services map[string][]string) {
+	serviceMap := make(map[consulService]bool)
+	for serviceName, tags := range services {
+		for _, tag := range tags {
+			service := consulService{
+				Tag:         tag,
+				ServiceName: serviceName,
+			}
+			serviceMap[service] = true
+			if _, ok := reg.watchingServices[service]; !ok {
+				err := reg.subscribe(serviceName)
+				if err != nil {
+					reg.logger.Errorf("failed to subscribe service, err: %v, service: %v", err, serviceName)
+					delete(serviceMap, service)
+				}
+			}
+		}
+	}
+
+	prevFetchServices := reg.watchingServices
+	reg.watchingServices = serviceMap
+
+	for key := range prevFetchServices {
+		if _, ok := serviceMap[key]; !ok {
+			err := reg.unsubscribe(key.ServiceName)
+			if err != nil {
+				reg.logger.Errorf("failed to unsubscribe service, err: %v, service: %v", err, key)
+			}
+			reg.softDeletedServices[key] = true
+		}
+	}
+
+}

controller/registries/consul/config_test.go

@@ -15,8 +15,12 @@
 package consul
 
 import (
+	"errors"
+	"reflect"
 	"testing"
 
+	"github.com/agiledragon/gomonkey/v2"
+	"github.com/hashicorp/consul/api"
 	"github.com/stretchr/testify/assert"
 
 	"mosn.io/htnn/controller/pkg/registry/log"
@@ -50,28 +54,42 @@ func TestStart(t *testing.T) {
 		logger: log.NewLogger(&log.RegistryLoggerOptions{
 			Name: "test",
 		}),
-		softDeletedServices: map[consulService]bool{},
-		done:                make(chan struct{}),
-		watchingServices:    map[consulService]bool{},
-	}
-	config := &consul.Config{
-		ServerUrl: "http://127.0.0.1:8500",
+		done: make(chan struct{}),
 	}
 
+	patches := gomonkey.ApplyPrivateMethod(reflect.TypeOf(reg), "fetchAllServices", func(_ *Consul, client *Client) (map[consulService]bool, error) {
+		return map[consulService]bool{
+			{ServiceName: "service1", Tag: "tag1"}: true,
+			{ServiceName: "service2", Tag: "tag2"}: true,
+		}, nil
+	})
+	config := &consul.Config{}
 	err := reg.Start(config)
-	assert.NoError(t, err)
-
+	assert.Nil(t, err)
 	err = reg.subscribe("123")
 	assert.Nil(t, err)
 
 	err = reg.unsubscribe("123")
 	assert.Nil(t, err)
 
-	err = reg.refresh()
-	assert.Nil(t, err)
-
 	err = reg.Stop()
 	assert.Nil(t, err)
+
+	patches.Reset()
+
+	config = &consul.Config{}
+
+	reg = &Consul{
+		logger: log.NewLogger(&log.RegistryLoggerOptions{
+			Name: "test",
+		}),
+		done: make(chan struct{}),
+	}
+
+	err = reg.Start(config)
+	assert.Error(t, err)
+
+	close(reg.done)
 }
 
 func TestReload(t *testing.T) {
@@ -83,3 +101,105 @@ func TestReload(t *testing.T) {
 	err := reg.Reload(config)
 	assert.NoError(t, err)
 }
+
+func TestRefresh(t *testing.T) {
+	reg := &Consul{
+		logger: log.NewLogger(&log.RegistryLoggerOptions{
+			Name: "test",
+		}),
+		softDeletedServices: map[consulService]bool{},
+		done:                make(chan struct{}),
+		watchingServices:    map[consulService]bool{},
+	}
+
+	config := &consul.Config{
+		ServerUrl: "http://127.0.0.1:8500",
+	}
+	client, _ := reg.NewClient(config)
+	reg.client = client
+	services := map[string][]string{
+		"service1": {"dc1", "dc2"},
+		"service2": {"dc1"},
+	}
+
+	reg.refresh(services)
+
+	assert.Len(t, reg.watchingServices, 3)
+	assert.Contains(t, reg.watchingServices, consulService{ServiceName: "service1", Tag: "dc1"})
+	assert.Contains(t, reg.watchingServices, consulService{ServiceName: "service1", Tag: "dc2"})
+	assert.Contains(t, reg.watchingServices, consulService{ServiceName: "service2", Tag: "dc1"})
+	assert.Empty(t, reg.softDeletedServices)
+
+	reg = &Consul{
+		logger: log.NewLogger(&log.RegistryLoggerOptions{
+			Name: "test",
+		}),
+		softDeletedServices: map[consulService]bool{},
+		watchingServices: map[consulService]bool{
+			{ServiceName: "service1", Tag: "dc1"}: true,
+		},
+	}
+
+	services = map[string][]string{}
+
+	reg.refresh(services)
+
+	assert.Len(t, reg.watchingServices, 0)
+	assert.Len(t, reg.softDeletedServices, 1)
+
+}
+
+func TestFetchAllServices(t *testing.T) {
+	t.Run("Test fetchAllServices method", func(t *testing.T) {
+		reg := &Consul{
+			logger: log.NewLogger(&log.RegistryLoggerOptions{
+				Name: "test",
+			}),
+		}
+		client := &Client{
+			consulCatalog: &api.Catalog{},
+			DataCenter:    "dc1",
+			NameSpace:     "ns1",
+			Token:         "token",
+		}
+
+		patches := gomonkey.ApplyMethod(reflect.TypeOf(client.consulCatalog), "Services", func(_ *api.Catalog, q *api.QueryOptions) (map[string][]string, *api.QueryMeta, error) {
+			return map[string][]string{
+				"service1": {"tag1", "tag2"},
+				"service2": {"tag3"},
+			}, nil, nil
+		})
+		defer patches.Reset()
+
+		services, err := reg.fetchAllServices(client)
+		assert.NoError(t, err)
+		assert.NotNil(t, services)
+		assert.True(t, services[consulService{ServiceName: "service1", Tag: "tag1"}])
+		assert.True(t, services[consulService{ServiceName: "service1", Tag: "tag2"}])
+		assert.True(t, services[consulService{ServiceName: "service2", Tag: "tag3"}])
+	})
+
+	t.Run("Test fetchAllServices method with error", func(t *testing.T) {
+		reg := &Consul{
+			logger: log.NewLogger(&log.RegistryLoggerOptions{
+				Name: "test",
+			}),
+		}
+		client := &Client{
+			consulCatalog: &api.Catalog{},
+			DataCenter:    "dc1",
+			NameSpace:     "ns1",
+			Token:         "token",
+		}
+
+		patches := gomonkey.ApplyMethod(reflect.TypeOf(client.consulCatalog), "Services", func(_ *api.Catalog, q *api.QueryOptions) (map[string][]string, *api.QueryMeta, error) {
+			return nil, nil, errors.New("mock error")
+		})
+		defer patches.Reset()
+
+		services, err := reg.fetchAllServices(client)
+		assert.Error(t, err)
+		assert.Equal(t, "mock error", err.Error())
+		assert.Nil(t, services)
+	})
+}