fix: harden v6 client lifecycle

Author: zouyx

modern_cache.go

@@ -15,6 +15,7 @@
 package agollo
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -93,15 +94,21 @@ func (c *ApolloClient) persistLocalSnapshot(snapshot ConfigSnapshot) error {
 	return nil
 }
 
-func (c *ApolloClient) loadLocalSnapshot(key ConfigKey) (ConfigSnapshot, error) {
+func (c *ApolloClient) loadLocalSnapshot(ctx context.Context, key ConfigKey) (ConfigSnapshot, error) {
 	for _, file := range []string{c.cacheFile(key), c.legacyCacheFile(key)} {
+		if err := ctx.Err(); err != nil {
+			return ConfigSnapshot{}, err
+		}
 		body, err := os.ReadFile(file)
 		if err != nil {
 			if errors.Is(err, os.ErrNotExist) {
 				continue
 			}
 			return ConfigSnapshot{}, fmt.Errorf("agollo: read local cache: %w", err)
 		}
+		if err := ctx.Err(); err != nil {
+			return ConfigSnapshot{}, err
+		}
 		if snapshot, err := decodeDiskSnapshot(key, body); err == nil {
 			return snapshot, nil
 		}

modern_client.go

@@ -286,6 +286,9 @@ func (c *ApolloClient) Load(ctx context.Context, namespaces ...string) error {
 }
 
 func (c *ApolloClient) getState(ctx context.Context, appID, namespace string, format ConfigFileFormat) (*modernConfig, error) {
+	if ctx == nil {
+		return nil, errors.New("agollo: context is nil")
+	}
 	appID = strings.TrimSpace(appID)
 	namespace = strings.TrimSpace(namespace)
 	if appID == "" {
@@ -374,6 +377,31 @@ func (c *ApolloClient) goBackground(run func(context.Context)) {
 	}()
 }
 
+func (c *ApolloClient) beginSubscription() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.closed {
+		return false
+	}
+	c.wg.Add(1)
+	return true
+}
+
+// loadContext is canceled when either the operation context or the client
+// lifecycle context ends. The watcher exits as soon as the returned cancel
+// function is called, so individual loads do not leave goroutines behind.
+func (c *ApolloClient) loadContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	loadContext, cancel := context.WithCancel(ctx)
+	go func() {
+		select {
+		case <-c.ctx.Done():
+			cancel()
+		case <-loadContext.Done():
+		}
+	}()
+	return loadContext, cancel
+}
+
 func normalizeURLs(urls []string) []string {
 	result := make([]string, 0, len(urls))
 	for _, rawURL := range urls {

modern_client_test.go

@@ -17,8 +17,10 @@ package agollo
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"regexp"
 	"strings"
 	"sync"
@@ -465,6 +467,40 @@ func TestApolloClientCloseRejectsNewStateAndSubscriptions(t *testing.T) {
 	}
 }
 
+func TestApolloClientCloseWaitsForConfigSubscription(t *testing.T) {
+	client := newTestClient(t, "http://127.0.0.1", ClientOptions{AppID: "sample"})
+	state := newModernConfig(client, ConfigKey{AppID: "sample", Cluster: "default", Namespace: "application", Format: ConfigFileFormatProperties})
+	client.mu.Lock()
+	client.states[state.key] = state
+	client.mu.Unlock()
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+	state.Subscribe(func(ConfigChangeEvent) {
+		close(started)
+		<-release
+	})
+	state.publish(ConfigSnapshot{Values: map[string]interface{}{"key": "value"}, Source: ConfigSourceRemote})
+	<-started
+
+	closed := make(chan error, 1)
+	go func() { closed <- client.Close() }()
+	select {
+	case err := <-closed:
+		t.Fatalf("Close() returned before listener finished: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+	close(release)
+	select {
+	case err := <-closed:
+		if err != nil {
+			t.Fatalf("Close() error = %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Close() did not wait for listener completion")
+	}
+}
+
 func TestDecodeDiskSnapshotRejectsDifferentCluster(t *testing.T) {
 	t.Parallel()
 	body, err := json.Marshal(diskSnapshot{
@@ -590,6 +626,117 @@ func TestApolloClientOfflineLoadsConfigMapWithoutCacheDirectory(t *testing.T) {
 	}
 }
 
+func TestApolloClientOfflineConfigMapDoesNotReadWorkingDirectoryCache(t *testing.T) {
+	workingDirectory, err := os.Getwd()
+	if err != nil {
+		t.Fatalf("Getwd() error = %v", err)
+	}
+	temporaryDirectory := t.TempDir()
+	if err := os.Chdir(temporaryDirectory); err != nil {
+		t.Fatalf("Chdir() error = %v", err)
+	}
+	t.Cleanup(func() { _ = os.Chdir(workingDirectory) })
+	if err := os.WriteFile("sample-application.json", []byte(`{"appId":"sample","cluster":"default","namespaceName":"application","configurations":{"key":"disk"}}`), 0o600); err != nil {
+		t.Fatalf("WriteFile() error = %v", err)
+	}
+
+	store := &memoryConfigMapStore{snapshot: ConfigSnapshot{
+		Key:    ConfigKey{AppID: "sample", Cluster: "default", Namespace: "application", Format: ConfigFileFormatProperties},
+		Values: map[string]interface{}{"key": "configmap"},
+	}}
+	client, err := NewClient(context.Background(), ClientOptions{AppID: "sample", ConfigMapStore: store, Offline: true})
+	if err != nil {
+		t.Fatalf("NewClient() error = %v", err)
+	}
+	defer client.Close()
+	config, err := client.Config(context.Background(), "application")
+	if err != nil {
+		t.Fatalf("Config() error = %v", err)
+	}
+	if got := config.String("key", ""); got != "configmap" || config.Source() != ConfigSourceConfigMap {
+		t.Fatalf("offline ConfigMap value = %q from %s", got, config.Source())
+	}
+}
+
+func TestApolloClientLoadHonorsOperationAndLifecycleContexts(t *testing.T) {
+	t.Run("operation context cancels ConfigMap load", func(t *testing.T) {
+		store := &blockingConfigMapStore{started: make(chan struct{}, 1)}
+		client, err := NewClient(context.Background(), ClientOptions{AppID: "sample", ConfigMapStore: store, Offline: true})
+		if err != nil {
+			t.Fatalf("NewClient() error = %v", err)
+		}
+		defer client.Close()
+		ctx, cancel := context.WithCancel(context.Background())
+		result := make(chan error, 1)
+		go func() {
+			_, err := client.Config(ctx, "application")
+			result <- err
+		}()
+		<-store.started
+		cancel()
+		select {
+		case err := <-result:
+			if !errors.Is(err, context.Canceled) {
+				t.Fatalf("Config() error = %v, want context cancellation", err)
+			}
+		case <-time.After(time.Second):
+			t.Fatal("ConfigMap load did not observe operation cancellation")
+		}
+	})
+
+	t.Run("client lifecycle cancels remote load", func(t *testing.T) {
+		started := make(chan struct{}, 1)
+		server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) {
+			select {
+			case started <- struct{}{}:
+			default:
+			}
+			<-request.Context().Done()
+		}))
+		defer server.Close()
+		clientContext, cancelClient := context.WithCancel(context.Background())
+		client, err := NewClient(clientContext, ClientOptions{AppID: "sample", ConfigServices: []string{server.URL}, DisableLongPolling: true})
+		if err != nil {
+			t.Fatalf("NewClient() error = %v", err)
+		}
+		defer client.Close()
+		result := make(chan error, 1)
+		go func() {
+			_, err := client.Config(context.Background(), "application")
+			result <- err
+		}()
+		<-started
+		cancelClient()
+		select {
+		case err := <-result:
+			if !errors.Is(err, context.Canceled) {
+				t.Fatalf("Config() error = %v, want context cancellation", err)
+			}
+		case <-time.After(time.Second):
+			t.Fatal("remote load did not observe client cancellation")
+		}
+	})
+}
+
+func TestApolloClientIntSliceReadsNativeIntSlice(t *testing.T) {
+	store := &memoryConfigMapStore{snapshot: ConfigSnapshot{
+		Key:    ConfigKey{AppID: "sample", Cluster: "default", Namespace: "application", Format: ConfigFileFormatProperties},
+		Values: map[string]interface{}{"ports": []int{8080, 9090}},
+	}}
+	client, err := NewClient(context.Background(), ClientOptions{AppID: "sample", ConfigMapStore: store, Offline: true})
+	if err != nil {
+		t.Fatalf("NewClient() error = %v", err)
+	}
+	defer client.Close()
+	config, err := client.Config(context.Background(), "application")
+	if err != nil {
+		t.Fatalf("Config() error = %v", err)
+	}
+	if got := config.IntSlice("ports", nil); len(got) != 2 || got[0] != 8080 || got[1] != 9090 {
+		t.Fatalf("IntSlice(ports) = %v", got)
+	}
+}
+
 func TestApolloClientDiscoversConfigServiceFromMetaServer(t *testing.T) {
 	t.Parallel()
 	var server *httptest.Server
@@ -651,8 +798,29 @@ func TestApolloClientLongPollBuildsDataCenterAndRefreshes(t *testing.T) {
 		}
 	}))
 	defer server.Close()
+	wrongServer := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) {
+		t.Fatalf("selector was bypassed for %s", request.URL.Path)
+	}))
+	defer wrongServer.Close()
 
-	client := newTestClient(t, server.URL, ClientOptions{AppID: "sample", DataCenter: "sh", ClientIP: "127.0.0.2"})
+	var selectorCalls atomic.Int32
+	client, err := NewClient(context.Background(), ClientOptions{
+		AppID:              "sample",
+		DataCenter:         "sh",
+		ClientIP:           "127.0.0.2",
+		ConfigServices:     []string{wrongServer.URL, server.URL},
+		DisableLongPolling: true,
+		ConfigServiceSelector: func(appID string, services []string) (string, error) {
+			if appID != "sample" || len(services) != 2 {
+				t.Fatalf("selector input = %q, %v", appID, services)
+			}
+			selectorCalls.Add(1)
+			return server.URL, nil
+		},
+	})
+	if err != nil {
+		t.Fatalf("NewClient() error = %v", err)
+	}
 	defer client.Close()
 	config, err := client.Config(context.Background(), "application")
 	if err != nil {
@@ -665,6 +833,9 @@ func TestApolloClientLongPollBuildsDataCenterAndRefreshes(t *testing.T) {
 	if got := config.String("key", ""); got != "after" {
 		t.Fatalf("long-poll refreshed value = %q", got)
 	}
+	if got := selectorCalls.Load(); got != 3 {
+		t.Fatalf("selector calls = %d, want initial fetch, long-poll, and refresh selection", got)
+	}
 }
 
 func TestApolloClientCloseCancelsLongPoll(t *testing.T) {
@@ -761,6 +932,21 @@ type memoryConfigMapStore struct {
 	saved    chan ConfigSnapshot
 }
 
+type blockingConfigMapStore struct {
+	started chan struct{}
+}
+
+func (s *blockingConfigMapStore) Load(ctx context.Context, _ ConfigKey) (ConfigSnapshot, error) {
+	select {
+	case s.started <- struct{}{}:
+	default:
+	}
+	<-ctx.Done()
+	return ConfigSnapshot{}, ctx.Err()
+}
+
+func (s *blockingConfigMapStore) Save(context.Context, ConfigSnapshot) error { return nil }
+
 func (s *memoryConfigMapStore) Load(ctx context.Context, key ConfigKey) (ConfigSnapshot, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()

modern_config.go

@@ -187,9 +187,13 @@ func (c *modernConfig) Subscribe(listener ConfigChangeHandler, options ...Subscr
 		c.listenerMu.Unlock()
 		return func() {}
 	}
+	if !c.client.beginSubscription() {
+		c.listenerMu.Unlock()
+		return func() {}
+	}
 	id := c.nextListener
 	c.nextListener++
-	subscription := newConfigSubscription(listener, settings, c.client.options.listenerQueueSize, &c.client.monitor)
+	subscription := newConfigSubscription(listener, settings, c.client.options.listenerQueueSize, &c.client.monitor, &c.client.wg)
 	c.listeners[id] = subscription
 	c.listenerMu.Unlock()
 
@@ -367,9 +371,13 @@ func (c *modernConfigFile) Subscribe(listener ConfigFileChangeHandler) func() {
 		config.listenerMu.Unlock()
 		return func() {}
 	}
+	if !config.client.beginSubscription() {
+		config.listenerMu.Unlock()
+		return func() {}
+	}
 	id := config.nextListener
 	config.nextListener++
-	subscription := newFileSubscription(listener, config.client.options.listenerQueueSize, &config.client.monitor)
+	subscription := newFileSubscription(listener, config.client.options.listenerQueueSize, &config.client.monitor, &config.client.wg)
 	config.fileListeners[id] = subscription
 	config.listenerMu.Unlock()
 	return func() {
@@ -425,10 +433,11 @@ type configSubscription struct {
 	done     chan struct{}
 	once     sync.Once
 	monitor  *modernMonitor
+	wg       *sync.WaitGroup
 }
 
-func newConfigSubscription(listener ConfigChangeHandler, options subscribeOptions, size int, monitor *modernMonitor) *configSubscription {
-	subscription := &configSubscription{listener: listener, options: options, queue: make(chan ConfigChangeEvent, size), done: make(chan struct{}), monitor: monitor}
+func newConfigSubscription(listener ConfigChangeHandler, options subscribeOptions, size int, monitor *modernMonitor, wg *sync.WaitGroup) *configSubscription {
+	subscription := &configSubscription{listener: listener, options: options, queue: make(chan ConfigChangeEvent, size), done: make(chan struct{}), monitor: monitor, wg: wg}
 	go subscription.run()
 	return subscription
 }
@@ -458,6 +467,7 @@ func (s *configSubscription) offer(event ConfigChangeEvent) {
 }
 
 func (s *configSubscription) run() {
+	defer s.wg.Done()
 	for {
 		select {
 		case <-s.done:
@@ -479,10 +489,11 @@ type fileSubscription struct {
 	done     chan struct{}
 	once     sync.Once
 	monitor  *modernMonitor
+	wg       *sync.WaitGroup
 }
 
-func newFileSubscription(listener ConfigFileChangeHandler, size int, monitor *modernMonitor) *fileSubscription {
-	subscription := &fileSubscription{listener: listener, queue: make(chan ConfigFileChangeEvent, size), done: make(chan struct{}), monitor: monitor}
+func newFileSubscription(listener ConfigFileChangeHandler, size int, monitor *modernMonitor, wg *sync.WaitGroup) *fileSubscription {
+	subscription := &fileSubscription{listener: listener, queue: make(chan ConfigFileChangeEvent, size), done: make(chan struct{}), monitor: monitor, wg: wg}
 	go subscription.run()
 	return subscription
 }
@@ -510,6 +521,7 @@ func (s *fileSubscription) offer(event ConfigFileChangeEvent) {
 }
 
 func (s *fileSubscription) run() {
+	defer s.wg.Done()
 	for {
 		select {
 		case <-s.done:

modern_poller.go

@@ -80,8 +80,14 @@ func (p *appPoller) poll() error {
 	if err != nil {
 		return err
 	}
-	serviceURL := p.client.nextConfigService(p.appID)
-	if serviceURL == "" || len(services) == 0 {
+	if len(services) == 0 {
+		return errors.New("agollo: no Config Service for long poll")
+	}
+	serviceURL, err := p.client.selectConfigService(p.appID, services)
+	if err != nil {
+		return err
+	}
+	if serviceURL == "" {
 		return errors.New("agollo: no Config Service for long poll")
 	}
 	endpoint, err := p.notificationsURL(serviceURL, states)