diff --git a/cli/server.go b/cli/server.go index e8b8768eea..758369de30 100644 --- a/cli/server.go +++ b/cli/server.go @@ -7,6 +7,7 @@ import ( "crypto/ecdsa" "crypto/elliptic" "crypto/rand" + "crypto/sha256" "crypto/tls" "crypto/x509" "database/sql" @@ -97,6 +98,7 @@ import ( "github.com/coder/coder/v2/coderd/workspaceapps/appurl" "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/coderd/wsbuilder" + "github.com/coder/coder/v2/coderd/x/nats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpcsdk" "github.com/coder/coder/v2/cryptorand" @@ -777,16 +779,34 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. } options.Database = database.New(sqlDB) - ps, err := pubsub.New(ctx, logger.Named("pubsub"), sqlDB, dbURL) + experiments := coderd.ReadExperiments(options.Logger, options.DeploymentValues.Experiments.Value()) + + pgPubsub, err := pubsub.New(ctx, logger.Named("pubsub"), sqlDB, dbURL) if err != nil { return xerrors.Errorf("create pubsub: %w", err) } - options.Pubsub = ps + options.Pubsub = pgPubsub + options.ReplicaSyncPubsub = pgPubsub + defer pgPubsub.Close() + if options.DeploymentValues.Prometheus.Enable { - options.PrometheusRegistry.MustRegister(ps) + options.PrometheusRegistry.MustRegister(pgPubsub) } - defer options.Pubsub.Close() - psWatchdog := pubsub.NewWatchdog(ctx, logger.Named("pswatch"), ps) + + // Use NATS for pubsub if the experiment is enabled. + if experiments.Enabled(codersdk.ExperimentNATSPubsub) { + token := fmt.Sprintf("%x", sha256.Sum256([]byte(dbURL))) + natsps, err := nats.New(ctx, logger.Named("pubsub"), nats.Options{ + ClusterAuthToken: token, + }) + if err != nil { + return xerrors.Errorf("create nats pubsub: %w", err) + } + options.Pubsub = natsps + defer natsps.Close() + } + + psWatchdog := pubsub.NewWatchdog(ctx, logger.Named("pswatch"), options.Pubsub) pubsubWatchdogTimeout = psWatchdog.Timeout() defer psWatchdog.Close() diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index f4c6ce20fb..aab6699a95 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -19251,12 +19251,14 @@ const docTemplate = `{ "workspace-usage", "oauth2", "mcp-server-http", - "workspace-build-updates" + "workspace-build-updates", + "nats_pubsub" ], "x-enum-comments": { "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", + "ExperimentNATSPubsub": "Enables embedded NATS pubsub.", "ExperimentNotifications": "Sends notifications via SMTP and webhooks following certain events.", "ExperimentOAuth2": "Enables OAuth2 provider functionality.", "ExperimentWorkspaceBuildUpdates": "Enables publishing workspace build updates to the all builds pubsub channel.", @@ -19269,7 +19271,8 @@ const docTemplate = `{ "Enables the new workspace usage tracking.", "Enables OAuth2 provider functionality.", "Enables the MCP HTTP server functionality.", - "Enables publishing workspace build updates to the all builds pubsub channel." + "Enables publishing workspace build updates to the all builds pubsub channel.", + "Enables embedded NATS pubsub." ], "x-enum-varnames": [ "ExperimentExample", @@ -19278,7 +19281,8 @@ const docTemplate = `{ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", - "ExperimentWorkspaceBuildUpdates" + "ExperimentWorkspaceBuildUpdates", + "ExperimentNATSPubsub" ] }, "codersdk.ExternalAPIKeyScopes": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 029b0dec8c..865e9ac96d 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -17489,12 +17489,14 @@ "workspace-usage", "oauth2", "mcp-server-http", - "workspace-build-updates" + "workspace-build-updates", + "nats_pubsub" ], "x-enum-comments": { "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", + "ExperimentNATSPubsub": "Enables embedded NATS pubsub.", "ExperimentNotifications": "Sends notifications via SMTP and webhooks following certain events.", "ExperimentOAuth2": "Enables OAuth2 provider functionality.", "ExperimentWorkspaceBuildUpdates": "Enables publishing workspace build updates to the all builds pubsub channel.", @@ -17507,7 +17509,8 @@ "Enables the new workspace usage tracking.", "Enables OAuth2 provider functionality.", "Enables the MCP HTTP server functionality.", - "Enables publishing workspace build updates to the all builds pubsub channel." + "Enables publishing workspace build updates to the all builds pubsub channel.", + "Enables embedded NATS pubsub." ], "x-enum-varnames": [ "ExperimentExample", @@ -17516,7 +17519,8 @@ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", - "ExperimentWorkspaceBuildUpdates" + "ExperimentWorkspaceBuildUpdates", + "ExperimentNATSPubsub" ] }, "codersdk.ExternalAPIKeyScopes": { diff --git a/coderd/coderd.go b/coderd/coderd.go index 6d8fa52208..875fad3c05 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -163,7 +163,10 @@ type Options struct { Logger slog.Logger Database database.Store Pubsub pubsub.Pubsub - RuntimeConfig *runtimeconfig.Manager + // ReplicaSyncPubsub is used explicitly to instantiate the replicasync manager downstream if it exists. + // All other consumers of pubsub should reference Options.Pubsub. + ReplicaSyncPubsub *pubsub.PGPubsub + RuntimeConfig *runtimeconfig.Manager // CacheDir is used for caching files served by the API. CacheDir string @@ -619,10 +622,9 @@ func New(options *Options) *API { ctx: ctx, cancel: cancel, DeploymentID: depID, - - ID: uuid.New(), - Options: options, - RootHandler: r, + ID: uuid.New(), + Options: options, + RootHandler: r, HTTPAuth: &HTTPAuthorizer{ Authorizer: options.Authorizer, Logger: options.Logger, diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index ab8d2271c3..94c6fde72f 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -166,8 +166,9 @@ type Options struct { // Overriding the database is heavily discouraged. // It should only be used in cases where multiple Coder // test instances are running against the same database. - Database database.Store - Pubsub pubsub.Pubsub + Database database.Store + Pubsub pubsub.Pubsub + ReplicaSyncPubsub *pubsub.PGPubsub // APIMiddleware inserts middleware before api.RootHandler, this can be // useful in certain tests where you want to intercept requests before @@ -287,6 +288,11 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can if options.Database == nil { options.Database, options.Pubsub = dbtestutil.NewDB(t) } + if options.ReplicaSyncPubsub == nil { + pgPubsub, ok := options.Pubsub.(*pubsub.PGPubsub) + require.True(t, ok, "ReplicaSyncPubsub must be a PGPubsub") + options.ReplicaSyncPubsub = pgPubsub + } if options.CoordinatorResumeTokenProvider == nil { options.CoordinatorResumeTokenProvider = tailnet.NewInsecureTestResumeTokenProvider() } @@ -596,6 +602,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can RuntimeConfig: runtimeManager, Database: options.Database, Pubsub: options.Pubsub, + ReplicaSyncPubsub: options.ReplicaSyncPubsub, ExternalAuthConfigs: options.ExternalAuthConfigs, UsageInserter: usageInserter, diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 9b36d11c27..b6e959b294 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -3415,8 +3415,9 @@ func TestReinit(t *testing.T) { triedToSubscribe: make(chan string), } client := coderdtest.New(t, &coderdtest.Options{ - Database: db, - Pubsub: &pubsubSpy, + Database: db, + Pubsub: &pubsubSpy, + ReplicaSyncPubsub: ps.(*pubsub.PGPubsub), }) user := coderdtest.CreateFirstUser(t, client) diff --git a/coderd/workspacestats/tracker_test.go b/coderd/workspacestats/tracker_test.go index fde8c9f2da..1ea81f63fb 100644 --- a/coderd/workspacestats/tracker_test.go +++ b/coderd/workspacestats/tracker_test.go @@ -113,11 +113,11 @@ func TestTracker_MultipleInstances(t *testing.T) { // Given we have two coderd instances connected to the same database var ( - ctx = testutil.Context(t, testutil.WaitLong) - db, _ = dbtestutil.NewDB(t) + ctx = testutil.Context(t, testutil.WaitLong) + db, ps = dbtestutil.NewDB(t) // real pubsub is not safe for concurrent use, and this test currently // does not depend on pubsub - ps = pubsub.NewInMemory() + psmem = pubsub.NewInMemory() wuTickA = make(chan time.Time) wuFlushA = make(chan int, 1) wuTickB = make(chan time.Time) @@ -132,7 +132,8 @@ func TestTracker_MultipleInstances(t *testing.T) { WorkspaceUsageTrackerTick: wuTickB, WorkspaceUsageTrackerFlush: wuFlushB, Database: db, - Pubsub: ps, + Pubsub: psmem, + ReplicaSyncPubsub: ps.(*pubsub.PGPubsub), }) owner = coderdtest.CreateFirstUser(t, clientA) now = dbtime.Now() diff --git a/coderd/x/nats/cluster.go b/coderd/x/nats/cluster.go index 7b0fd1ab80..aa12c748fe 100644 --- a/coderd/x/nats/cluster.go +++ b/coderd/x/nats/cluster.go @@ -1,6 +1,7 @@ package nats import ( + "errors" "net" "net/url" "slices" @@ -8,10 +9,68 @@ import ( "strings" "golang.org/x/xerrors" + + "cdr.dev/slog/v3" ) -// SetPeerAddresses replaces the configured NATS cluster peer routes. -func (p *Pubsub) SetPeerAddresses(addresses []string) error { +const defaultClusterTokenUsername = "coder" + +// PeerFetcher fetches NATS peer route addresses. +type PeerFetcher interface { + PrimaryPeerAddresses() []string +} + +type NopPeerFetcher struct{} + +func (NopPeerFetcher) PrimaryPeerAddresses() []string { + return nil +} + +// SetPeerFetcher replaces the peer fetcher used by RefreshPeers and triggers +// an immediate peer refresh. Passing nil disables peering. +func (p *Pubsub) SetPeerFetcher(fetcher PeerFetcher) { + p.mu.Lock() + if fetcher == nil { + fetcher = NopPeerFetcher{} + } + p.peerFetcher = fetcher + p.mu.Unlock() + p.RefreshPeers() +} + +// RefreshPeers signals the peer refresh worker to fetch and apply the latest +// peer route addresses. Multiple pending refreshes are coalesced. +func (p *Pubsub) RefreshPeers() { + select { + case p.peerRefresh <- struct{}{}: + default: + } +} + +func (p *Pubsub) runPeerRefresh() { + for { + p.mu.Lock() + fetcher := p.peerFetcher + p.mu.Unlock() + + addrs := fetcher.PrimaryPeerAddresses() + if err := p.setPeerAddresses(addrs); err != nil { + if errors.Is(err, errClosed) && p.ctx.Err() != nil { + return + } + p.logger.Error(p.ctx, "refresh nats peers", slog.Error(err)) + } + + select { + case <-p.ctx.Done(): + return + case <-p.peerRefresh: + } + } +} + +// setPeerAddresses replaces the configured NATS cluster peer routes. +func (p *Pubsub) setPeerAddresses(addresses []string) error { p.clusterMu.Lock() defer p.clusterMu.Unlock() @@ -22,13 +81,18 @@ func (p *Pubsub) SetPeerAddresses(addresses []string) error { return xerrors.New("nats pubsub was not started with clustering enabled") } - routes, err := parsePeerAddresses(addresses) + routes, err := p.parsePeerAddresses(addresses) if err != nil { return err } - self := &url.URL{Scheme: "nats", Host: p.ns.ClusterAddr().String()} + self := &url.URL{Scheme: "nats", Host: p.Server.ClusterAddr().String()} routes = filterSelfRoutes(routes, self) + + if p.opts.ClusterAuthToken != "" { + routes = routesWithAuth(routes, p.opts.ClusterAuthToken) + } + routes = sortRouteURLs(routes) if sortedURLsEqual(p.currentRoutes, routes) { @@ -37,7 +101,7 @@ func (p *Pubsub) SetPeerAddresses(addresses []string) error { newOpts := p.serverOpts.Clone() newOpts.Routes = cloneRouteURLs(routes) - if err := p.ns.ReloadOptions(newOpts); err != nil { + if err := p.Server.ReloadOptions(newOpts); err != nil { return xerrors.Errorf("reload nats peer addresses: %w", err) } p.serverOpts = newOpts.Clone() @@ -45,7 +109,7 @@ func (p *Pubsub) SetPeerAddresses(addresses []string) error { return nil } -func parsePeerAddresses(addresses []string) ([]*url.URL, error) { +func (p *Pubsub) parsePeerAddresses(addresses []string) ([]*url.URL, error) { routesByAddress := make(map[string]*url.URL, len(addresses)) for i, address := range addresses { trimmed := strings.TrimSpace(address) @@ -53,14 +117,25 @@ func parsePeerAddresses(addresses []string) ([]*url.URL, error) { return nil, xerrors.Errorf("peer address %d is empty", i) } - normalizedHost, err := normalizeHostPort(trimmed) + host, port, err := normalizeHostPort(trimmed) if err != nil { return nil, err } - routesByAddress[normalizedHost] = &url.URL{ + // This is a hack to enable testing with an arbitrary port. The logic here + // is to presume if the default port is being used then we are running in prod + // and all peers are using the same port. If the port is not the default then + // we are running a test in which case we should pass through the custom port. + // This hack will be removed when https://github.com/coder/scaletest/issues/149 + // is resolved. + if p.opts.ClusterPort == defaultClusterPort { + port = defaultClusterPort + } + + hostPort := net.JoinHostPort(host, strconv.Itoa(port)) + routesByAddress[hostPort] = &url.URL{ Scheme: "nats", - Host: normalizedHost, + Host: hostPort, } } @@ -82,34 +157,34 @@ func filterSelfRoutes(routes []*url.URL, self *url.URL) []*url.URL { return filtered } -func normalizeHostPort(address string) (string, error) { +func normalizeHostPort(address string) (string, int, error) { route, err := url.Parse(address) if err != nil { - return "", xerrors.Errorf("parse peer address %q: %w", address, err) + return "", 0, xerrors.Errorf("parse peer address %q: %w", address, err) } if route.User != nil { - return "", xerrors.Errorf("peer address %q must not include userinfo", address) + return "", 0, xerrors.Errorf("peer address %q must not include userinfo", address) } if route.Path != "" || route.RawQuery != "" || route.Fragment != "" { - return "", xerrors.Errorf("peer address %q must not include path, query, or fragment", address) + return "", 0, xerrors.Errorf("peer address %q must not include path, query, or fragment", address) } host, port, err := net.SplitHostPort(route.Host) if err != nil { - return "", xerrors.Errorf("split %q host port: %w", address, err) + return "", 0, xerrors.Errorf("split %q host port: %w", address, err) } if host == "" || port == "" { - return "", xerrors.Errorf("%q must include host and port", address) + return "", 0, xerrors.Errorf("%q must include host and port", address) } portNumber, err := strconv.Atoi(port) if err != nil { - return "", xerrors.Errorf("parse %q port: %w", address, err) + return "", 0, xerrors.Errorf("parse %q port: %w", address, err) } if portNumber <= 0 || portNumber > 65535 { - return "", xerrors.Errorf("peer address %q must include a valid port", address) + return "", 0, xerrors.Errorf("peer address %q must include a valid port", address) } - return net.JoinHostPort(host, strconv.Itoa(portNumber)), nil + return host, portNumber, nil } func sortRouteURLs(routes []*url.URL) []*url.URL { @@ -119,6 +194,23 @@ func sortRouteURLs(routes []*url.URL) []*url.URL { return routes } +func routesWithAuth(routes []*url.URL, token string) []*url.URL { + if token == "" { + return routes + } + withAuth := make([]*url.URL, 0, len(routes)) + for _, route := range routes { + if route == nil { + withAuth = append(withAuth, nil) + continue + } + clone := *route + clone.User = url.UserPassword(defaultClusterTokenUsername, token) + withAuth = append(withAuth, &clone) + } + return withAuth +} + // sortedURLsEqual assumes sorted slices. func sortedURLsEqual(a, b []*url.URL) bool { if len(a) != len(b) { diff --git a/coderd/x/nats/cluster_internal_test.go b/coderd/x/nats/cluster_internal_test.go index eadf2e561f..5d70d74f87 100644 --- a/coderd/x/nats/cluster_internal_test.go +++ b/coderd/x/nats/cluster_internal_test.go @@ -6,6 +6,8 @@ import ( "testing" "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/testutil" ) func Test_parsePeerAddresses(t *testing.T) { @@ -13,7 +15,8 @@ func Test_parsePeerAddresses(t *testing.T) { t.Run("Valid", func(t *testing.T) { t.Parallel() - routes, err := parsePeerAddresses([]string{ + ps := &Pubsub{} + routes, err := ps.parsePeerAddresses([]string{ "whatever://127.0.0.1:4222 ", "http://[::1]:7222", "nats://example.com:6222", @@ -26,16 +29,37 @@ func Test_parsePeerAddresses(t *testing.T) { }, routeStrings(routes)) }) + // Test that when a pubsub is running with the default port, it assumes all peers are also using + // the default port. + t.Run("PrefersDefaultPort", func(t *testing.T) { + t.Parallel() + ps := &Pubsub{} + ps.opts.ClusterPort = defaultClusterPort + routes, err := ps.parsePeerAddresses([]string{ + "whatever://127.0.0.1:4222 ", + "http://[::1]:7222", + "nats://example.com:1234", + }) + require.NoError(t, err) + require.ElementsMatch(t, []string{ + "nats://127.0.0.1:6222", + "nats://[::1]:6222", + "nats://example.com:6222", + }, routeStrings(routes)) + }) + t.Run("Empty", func(t *testing.T) { t.Parallel() - routes, err := parsePeerAddresses(nil) + ps := &Pubsub{} + routes, err := ps.parsePeerAddresses(nil) require.NoError(t, err) require.Empty(t, routes) }) t.Run("Dedupes", func(t *testing.T) { t.Parallel() - routes, err := parsePeerAddresses([]string{ + ps := &Pubsub{} + routes, err := ps.parsePeerAddresses([]string{ "nats://b.example:6222", "nats://a.example:6222", "nats://b.example:6222", @@ -68,7 +92,8 @@ func Test_parsePeerAddresses(t *testing.T) { } { t.Run(address, func(t *testing.T) { t.Parallel() - _, err := parsePeerAddresses([]string{address}) + ps := &Pubsub{} + _, err := ps.parsePeerAddresses([]string{address}) require.Error(t, err) }) } @@ -78,7 +103,8 @@ func Test_parsePeerAddresses(t *testing.T) { func Test_filterSelfRoutes(t *testing.T) { t.Parallel() - routes, err := parsePeerAddresses([]string{ + ps := &Pubsub{} + routes, err := ps.parsePeerAddresses([]string{ "nats://b.example:6222", "http://self.example:6222", }) @@ -88,24 +114,102 @@ func Test_filterSelfRoutes(t *testing.T) { require.Equal(t, []string{"nats://b.example:6222"}, routeStrings(routes)) } -// Cluster tests bind free ports and reload shared route state. -func TestPubsub_SetPeerAddresses(t *testing.T) { +func TestPubsub_RefreshPeers(t *testing.T) { + t.Parallel() + + t.Run("PeersFetchedOnStartup", func(t *testing.T) { + t.Parallel() + + // Supplying PeerFetcher in Options should be enough to seed routes. + // Callers should not need a separate SetPeerFetcher or RefreshPeers call + // after New returns. + fetcher := &testPeerFetcher{addresses: []string{"nats://127.0.0.1:1234"}} + opts := clusterTestOptions(t) + opts.PeerFetcher = fetcher + a := newTestPubsub(t, opts) + + require.Eventually(t, func() bool { + routes := currentRouteURLs(a) + return sortedURLsEqual(routes, sortRouteURLs(mustParsePeerAddresses(t, + addrWithAuth(t, "nats://127.0.0.1:1234", opts.ClusterAuthToken), + ))) + }, testutil.WaitShort, testutil.IntervalFast) + }) + + t.Run("SetPeerFetcher", func(t *testing.T) { + t.Parallel() + opts := clusterTestOptions(t) + a := newTestPubsub(t, opts) + + routes := []string{ + "nats://127.0.0.1:1234", + "nats://127.0.0.1:1235", + } + fetcher := &testPeerFetcher{routes} + + expectedRoutes := routesWithAuth(mustParsePeerAddresses(t, fetcher.addresses...), opts.ClusterAuthToken) + + a.SetPeerFetcher(fetcher) + require.Eventually(t, func() bool { + return sortedURLsEqual(currentRouteURLs(a), sortRouteURLs(expectedRoutes)) + }, testutil.WaitShort, testutil.IntervalFast) + + a.SetPeerFetcher(nil) + require.Eventually(t, func() bool { + return sortedURLsEqual(currentRouteURLs(a), nil) + }, testutil.WaitShort, testutil.IntervalFast) + }) +} + +func mustParsePeerAddresses(t *testing.T, addresses ...string) []*url.URL { + t.Helper() + routes := make([]*url.URL, 0, len(addresses)) + for _, address := range addresses { + route, err := url.Parse(address) + require.NoError(t, err) + routes = append(routes, route) + } + return routes +} + +func currentRouteURLs(ps *Pubsub) []*url.URL { + ps.clusterMu.Lock() + defer ps.clusterMu.Unlock() + return cloneRouteURLs(ps.currentRoutes) +} + +type testPeerFetcher struct { + addresses []string +} + +func (f *testPeerFetcher) PrimaryPeerAddresses() []string { + return f.addresses +} + +func TestPubsub_setPeerAddresses(t *testing.T) { t.Parallel() t.Run("OK", func(t *testing.T) { t.Parallel() - a := newTestPubsub(t, clusterTestOptions(t)) - b := newTestPubsub(t, clusterTestOptions(t)) - c := newTestPubsub(t, clusterTestOptions(t)) + opts := clusterTestOptions(t) + a := newTestPubsub(t, opts) + b := newTestPubsub(t, opts) + c := newTestPubsub(t, opts) addrB := clusterRouteAddress(t, b) addrC := clusterRouteAddress(t, c) - require.NoError(t, a.SetPeerAddresses([]string{addrC, addrB})) - requireRoutesEqual(t, a.currentRoutes, addrB, addrC) + require.NoError(t, a.setPeerAddresses([]string{addrC, addrB})) + requireRoutesEqual(t, a.currentRoutes, + addrWithAuth(t, addrB, opts.ClusterAuthToken), + addrWithAuth(t, addrC, opts.ClusterAuthToken), + ) - require.NoError(t, a.SetPeerAddresses([]string{addrB, addrC})) - requireRoutesEqual(t, a.currentRoutes, addrB, addrC) + require.NoError(t, a.setPeerAddresses([]string{addrB, addrC})) + requireRoutesEqual(t, a.currentRoutes, + addrWithAuth(t, addrB, opts.ClusterAuthToken), + addrWithAuth(t, addrC, opts.ClusterAuthToken), + ) - require.NoError(t, a.SetPeerAddresses(nil)) + require.NoError(t, a.setPeerAddresses(nil)) require.Empty(t, a.currentRoutes) require.Empty(t, a.serverOpts.Routes) }) @@ -113,7 +217,7 @@ func TestPubsub_SetPeerAddresses(t *testing.T) { t.Run("StandaloneConfigError", func(t *testing.T) { t.Parallel() ps := newTestPubsub(t, defaultTestOptions()) - err := ps.SetPeerAddresses(nil) + err := ps.setPeerAddresses(nil) require.ErrorContains(t, err, "not started with clustering enabled") }) @@ -121,14 +225,14 @@ func TestPubsub_SetPeerAddresses(t *testing.T) { t.Parallel() ps := newTestPubsub(t, clusterTestOptions(t)) require.NoError(t, ps.Close()) - err := ps.SetPeerAddresses(nil) + err := ps.setPeerAddresses(nil) require.True(t, errors.Is(err, errClosed), "got %v", err) }) t.Run("DropsSelfRoute", func(t *testing.T) { t.Parallel() ps := newTestPubsub(t, clusterTestOptions(t)) - require.NoError(t, ps.SetPeerAddresses([]string{clusterRouteAddress(t, ps)})) + require.NoError(t, ps.setPeerAddresses([]string{clusterRouteAddress(t, ps)})) require.Empty(t, ps.currentRoutes) }) } diff --git a/coderd/x/nats/pubsub.go b/coderd/x/nats/pubsub.go index a41247ed09..4c6d902fd2 100644 --- a/coderd/x/nats/pubsub.go +++ b/coderd/x/nats/pubsub.go @@ -81,6 +81,14 @@ type Options struct { // 6222 when cluster mode is enabled. ClusterPort int + // ClusterAuthToken is the shared route authentication token for + // clustered embedded NATS servers. Empty disables route auth. + ClusterAuthToken string + + // PeerFetcher provides the current set of peer route addresses. + // RefreshPeers uses it to update the configured cluster routes. + PeerFetcher PeerFetcher + // RoutePoolSize is the NATS route pool size. Zero means the package // default when cluster mode is enabled. RoutePoolSize int @@ -106,7 +114,7 @@ type Pubsub struct { logger slog.Logger opts Options - ns *natsserver.Server + Server *natsserver.Server // publishPool and subscribePool are immutable after construction so // the hot path can index without holding p.mu. publishPool []*natsgo.Conn @@ -126,6 +134,9 @@ type Pubsub struct { clustered bool serverOpts *natsserver.Options currentRoutes []*url.URL + + peerFetcher PeerFetcher + peerRefresh chan struct{} } // natsSub maps to one underlying *natsgo.Subscription. The first @@ -183,6 +194,8 @@ func newPubsub(ctx context.Context, logger slog.Logger, opts Options) *Pubsub { subscriptions: make(map[string]*natsSub), ctx: ctx, cancel: cancel, + peerFetcher: opts.PeerFetcher, + peerRefresh: make(chan struct{}, 1), } } @@ -246,8 +259,12 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Pubsub, error) slog.F("client_url", ns.ClientURL()), ) + if opts.PeerFetcher == nil { + opts.PeerFetcher = NopPeerFetcher{} + } + p := newPubsub(ctx, logger, opts) - p.ns = ns + p.Server = ns p.clustered = !opts.disableCluster p.serverOpts = sopts.Clone() p.currentRoutes = cloneRouteURLs(sopts.Routes) @@ -260,6 +277,7 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Pubsub, error) ns.WaitForShutdown() return nil, err } + subscribePool, err := newConnPool(ns, opts, handlers, opts.SubscribeConns, "coder-pubsub-sub") if err != nil { p.cancel() @@ -270,12 +288,18 @@ func New(ctx context.Context, logger slog.Logger, opts Options) (*Pubsub, error) ns.WaitForShutdown() return nil, err } + p.publishPool = publishPool p.subscribePool = subscribePool + + if p.clustered { + go p.runPeerRefresh() + } go func() { <-p.ctx.Done() _ = p.Close() }() + return p, nil } @@ -670,9 +694,9 @@ func (p *Pubsub) Close() error { } } - if p.ns != nil { - p.ns.Shutdown() - p.ns.WaitForShutdown() + if p.Server != nil { + p.Server.Shutdown() + p.Server.WaitForShutdown() } }) return nil diff --git a/coderd/x/nats/pubsub_internal_test.go b/coderd/x/nats/pubsub_internal_test.go index 3b5263654e..b6f55f046b 100644 --- a/coderd/x/nats/pubsub_internal_test.go +++ b/coderd/x/nats/pubsub_internal_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "net/url" + "slices" "sync" "sync/atomic" "testing" @@ -115,8 +116,8 @@ func Test_New(t *testing.T) { } }) - require.Equal(t, 2, ps.ns.NumClients(), - "expected exactly 2 client connections (pubConn + subConn), got %d", ps.ns.NumClients()) + require.Equal(t, 2, ps.Server.NumClients(), + "expected exactly 2 client connections (pubConn + subConn), got %d", ps.Server.NumClients()) require.Len(t, ps.publishPool, 1, "default PublishConns must be 1") require.Len(t, ps.subscribePool, 1, "default SubscribeConns must be 1") require.NotSame(t, ps.publishPool[0], ps.subscribePool[0], "pubConn and subConn must be distinct") @@ -329,13 +330,12 @@ func Test_localSub_init(t *testing.T) { require.Len(t, ps.subscribePool, 1) require.False(t, ps.subscribePool[0].IsClosed(), "subConn must not be closed by slow consumer") require.True(t, ps.subscribePool[0].IsConnected(), "subConn must stay connected") - require.Equal(t, 2, ps.ns.NumClients(), "slow consumer must not disconnect subConn") + require.Equal(t, 2, ps.Server.NumClients(), "slow consumer must not disconnect subConn") }) } func TestPubsubCluster(t *testing.T) { t.Parallel() - // OK verifies that SetPeerAddresses changes the active cluster topology. // A starts connected to B, then C is added and receives both global and // C-only messages. B is then removed from A's peers, while C continues to @@ -343,15 +343,18 @@ func TestPubsubCluster(t *testing.T) { t.Run("OK", func(t *testing.T) { t.Parallel() - a := newTestPubsub(t, clusterTestOptions(t)) - b := newTestPubsub(t, clusterTestOptions(t)) - c := newTestPubsub(t, clusterTestOptions(t)) + opts := clusterTestOptions(t) + a := newTestPubsub(t, opts) + b := newTestPubsub(t, opts) + c := newTestPubsub(t, opts) addrB := clusterRouteAddress(t, b) addrC := clusterRouteAddress(t, c) - require.NoError(t, a.SetPeerAddresses([]string{addrB})) - requireRoutesEqual(t, a.currentRoutes, addrB) + require.NoError(t, a.setPeerAddresses([]string{addrB})) + requireRoutesEqual(t, a.currentRoutes, + addrWithAuth(t, addrB, opts.ClusterAuthToken), + ) globalEvent := "global" bGlobal := make(chan []byte, 8) @@ -383,8 +386,11 @@ func TestPubsubCluster(t *testing.T) { // Add C to A's peer list. B and C should both receive global messages, // while the C-only subject should route only to C. - require.NoError(t, a.SetPeerAddresses([]string{addrC, addrB})) - requireRoutesEqual(t, a.currentRoutes, addrB, addrC) + require.NoError(t, a.setPeerAddresses([]string{addrC, addrB})) + requireRoutesEqual(t, a.currentRoutes, + addrWithAuth(t, addrB, opts.ClusterAuthToken), + addrWithAuth(t, addrC, opts.ClusterAuthToken), + ) waitForRouteSubscription(t, a, globalEvent) waitForRouteSubscription(t, a, cSubject) @@ -397,8 +403,10 @@ func TestPubsubCluster(t *testing.T) { require.Equal(t, "c-unique-msg", string(receiveMessage(t, cUnique))) // Remove B from A's peer list. Only C should receive the next messages. - require.NoError(t, a.SetPeerAddresses([]string{addrC})) - requireRoutesEqual(t, a.currentRoutes, addrC) + require.NoError(t, a.setPeerAddresses([]string{addrC})) + requireRoutesEqual(t, a.currentRoutes, + addrWithAuth(t, addrC, opts.ClusterAuthToken), + ) publishAndFlush(t, a, globalEvent, "no-b-peer") require.Equal(t, "no-b-peer", string(receiveMessage(t, cGlobal))) @@ -406,6 +414,60 @@ func TestPubsubCluster(t *testing.T) { publishAndFlush(t, a, cSubject, "c-messages-still-work") require.Equal(t, "c-messages-still-work", string(receiveMessage(t, cUnique))) }) + + // InvalidAuthRejected asserts the cluster route listener rejects + // connections that do not present the configured ClusterAuthToken. + // We dial the route listener directly with the nats.go client, which + // surfaces a typed nats.ErrAuthorization for protocol-level -ERR + // 'Authorization Violation' responses. + t.Run("ClusterAuthRequired", func(t *testing.T) { + t.Parallel() + + ps := newTestPubsub(t, clusterTestOptions(t)) + routeURL := clusterRouteAddress(t, ps) + + _, err := natsgo.Connect(routeURL, + natsgo.Token("wrong-token"), + natsgo.MaxReconnects(0), + natsgo.RetryOnFailedConnect(false), + natsgo.Timeout(testutil.WaitShort), + ) + require.ErrorIs(t, err, natsgo.ErrAuthorization, + "route dial with wrong token must be rejected") + + _, err = natsgo.Connect(routeURL, + natsgo.MaxReconnects(0), + natsgo.RetryOnFailedConnect(false), + natsgo.Timeout(testutil.WaitShort), + ) + require.ErrorIs(t, err, natsgo.ErrAuthorization, + "unauthenticated route dial must be rejected") + }) + + // ClientAuthRequired asserts the local NATS client listener also requires + // the configured ClusterAuthToken, so loopback clients cannot bypass auth. + t.Run("ClientAuthRequired", func(t *testing.T) { + t.Parallel() + + opts := clusterTestOptions(t) + ps := newTestPubsub(t, opts) + clientURL := ps.Server.ClientURL() + + _, err := natsgo.Connect(clientURL, + natsgo.MaxReconnects(0), + natsgo.RetryOnFailedConnect(false), + natsgo.Timeout(testutil.WaitShort), + ) + require.ErrorIs(t, err, natsgo.ErrAuthorization, + "unauthenticated client connect must be rejected") + + nc, err := natsgo.Connect(clientURL, + natsgo.Token(opts.ClusterAuthToken), + natsgo.Timeout(testutil.WaitShort), + ) + require.NoError(t, err, "authenticated client connect with matching token must succeed") + nc.Close() + }) } func defaultTestOptions() Options { @@ -415,9 +477,10 @@ func defaultTestOptions() Options { func clusterTestOptions(t *testing.T) Options { t.Helper() return Options{ - ClusterHost: "127.0.0.1", - ClusterPort: natsserver.RANDOM_PORT, - disableCluster: false, + ClusterHost: "127.0.0.1", + ClusterPort: natsserver.RANDOM_PORT, + disableCluster: false, + ClusterAuthToken: fmt.Sprintf("shared-token-%d", time.Now().UnixNano()), } } @@ -435,15 +498,23 @@ func newTestPubsub(t *testing.T, opts Options) *Pubsub { func clusterRouteAddress(t *testing.T, ps *Pubsub) string { t.Helper() - addr := ps.ns.ClusterAddr() + addr := ps.Server.ClusterAddr() require.NotNil(t, addr) return "nats://" + addr.String() } +func addrWithAuth(t *testing.T, addr string, authToken string) string { + t.Helper() + u, err := url.Parse(addr) + require.NoError(t, err) + u.User = url.UserPassword(defaultClusterTokenUsername, authToken) + return u.String() +} + func waitForRouteSubscription(t *testing.T, ps *Pubsub, subject string) { t.Helper() require.Eventually(t, func() bool { - routes, err := ps.ns.Routez(&natsserver.RoutezOptions{Subscriptions: true}) + routes, err := ps.Server.Routez(&natsserver.RoutezOptions{Subscriptions: true}) if err != nil { return false } @@ -477,16 +548,19 @@ func receiveMessage(t *testing.T, got <-chan []byte) []byte { func requireRoutesEqual(t *testing.T, routes []*url.URL, addresses ...string) { t.Helper() - want, err := parsePeerAddresses(addresses) - require.NoError(t, err) - want = sortRouteURLs(want) - require.True(t, sortedURLsEqual(want, routes), "want %v, got %v", routeStrings(want), routeStrings(routes)) + + rrs := routeStrings(routes) + + slices.Sort(rrs) + slices.Sort(addresses) + + require.True(t, slices.Equal(rrs, addresses), "want %v, got %v", rrs, addresses) } func routeStrings(routes []*url.URL) []string { - strings := make([]string, 0, len(routes)) + out := make([]string, 0, len(routes)) for _, route := range routes { - strings = append(strings, route.String()) + out = append(out, route.String()) } - return strings + return out } diff --git a/coderd/x/nats/server.go b/coderd/x/nats/server.go index 6013c44feb..47194c8a75 100644 --- a/coderd/x/nats/server.go +++ b/coderd/x/nats/server.go @@ -34,6 +34,9 @@ func buildServerOptions(opts Options) (*natsserver.Options, error) { sopts.DontListen = false sopts.Host = "127.0.0.1" sopts.Port = natsserver.RANDOM_PORT + if opts.ClusterAuthToken != "" { + sopts.Authorization = opts.ClusterAuthToken + } if !opts.disableCluster { clusterHost := opts.ClusterHost @@ -55,6 +58,10 @@ func buildServerOptions(opts Options) (*natsserver.Options, error) { Port: clusterPort, PoolSize: routePoolSize, } + if opts.ClusterAuthToken != "" { + sopts.Cluster.Username = defaultClusterTokenUsername + sopts.Cluster.Password = opts.ClusterAuthToken + } } return sopts, nil @@ -90,6 +97,9 @@ func connectClient(ns *natsserver.Server, opts Options, handlers connHandlers, c connOpts := []natsgo.Option{ natsgo.Name(connName), } + if opts.ClusterAuthToken != "" { + connOpts = append(connOpts, natsgo.Token(opts.ClusterAuthToken)) + } if opts.ReconnectWait > 0 { connOpts = append(connOpts, natsgo.ReconnectWait(opts.ReconnectWait)) } diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 8164222831..68d6ae0f7c 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -5003,6 +5003,7 @@ const ( ExperimentOAuth2 Experiment = "oauth2" // Enables OAuth2 provider functionality. ExperimentMCPServerHTTP Experiment = "mcp-server-http" // Enables the MCP HTTP server functionality. ExperimentWorkspaceBuildUpdates Experiment = "workspace-build-updates" // Enables publishing workspace build updates to the all builds pubsub channel. + ExperimentNATSPubsub Experiment = "nats_pubsub" // Enables embedded NATS pubsub. ) func (e Experiment) DisplayName() string { @@ -5021,6 +5022,8 @@ func (e Experiment) DisplayName() string { return "MCP HTTP Server Functionality" case ExperimentWorkspaceBuildUpdates: return "Workspace Build Updates Channel" + case ExperimentNATSPubsub: + return "NATS Pubsub" default: // Split on hyphen and convert to title case // e.g. "mcp-server-http" -> "Mcp Server Http" @@ -5037,6 +5040,7 @@ var ExperimentsKnown = Experiments{ ExperimentWorkspaceUsage, ExperimentOAuth2, ExperimentMCPServerHTTP, + ExperimentNATSPubsub, ExperimentWorkspaceBuildUpdates, } diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 91db75c177..55eac2c4f2 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7219,9 +7219,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o #### Enumerated Values -| Value(s) | -|-------------------------------------------------------------------------------------------------------------------------------| -| `auto-fill-parameters`, `example`, `mcp-server-http`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-usage` | +| Value(s) | +|----------------------------------------------------------------------------------------------------------------------------------------------| +| `auto-fill-parameters`, `example`, `mcp-server-http`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-usage` | ## codersdk.ExternalAPIKeyScopes diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 092314a973..2df327f674 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -45,6 +45,7 @@ import ( agplschedule "github.com/coder/coder/v2/coderd/schedule" agplusage "github.com/coder/coder/v2/coderd/usage" "github.com/coder/coder/v2/coderd/wsbuilder" + "github.com/coder/coder/v2/coderd/x/nats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/enterprise/aiseats" "github.com/coder/coder/v2/enterprise/coderd/connectionlog" @@ -655,7 +656,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { // We always want to run the replica manager even if we don't have DERP // enabled, since it's used to detect other coder servers for licensing. - api.replicaManager, err = replicasync.New(ctx, options.Logger, options.Database, options.Pubsub, &replicasync.Options{ + api.replicaManager, err = replicasync.New(ctx, options.Logger, options.Database, options.ReplicaSyncPubsub, &replicasync.Options{ ID: api.AGPL.ID, RelayAddress: options.DERPServerRelayAddress, // #nosec G115 - DERP region IDs are small and fit in int32 @@ -757,6 +758,10 @@ type Options struct { ExternalTokenEncryption []dbcrypt.Cipher + // ReplicaManager detects and syncs multiple Coder replicas. When provided, + // the API owns and closes it. + ReplicaManager *replicasync.Manager + // Used for high availability. ReplicaSyncUpdateInterval time.Duration ReplicaErrorGracePeriod time.Duration @@ -965,7 +970,12 @@ func (api *API) updateEntitlements(ctx context.Context) error { coordinator = haCoordinator } - api.replicaManager.SetCallback(func() { + if natsPubsub, ok := api.Pubsub.(*nats.Pubsub); ok { + natsPubsub.SetPeerFetcher(api.replicaManager) + api.replicaManager.SetCallback("nats", natsPubsub.RefreshPeers) + } + + api.replicaManager.SetCallback("derp", func() { // Only update DERP mesh if the built-in server is enabled. if api.Options.DeploymentValues.DERP.Server.Enable { addresses := make([]string, 0) @@ -985,11 +995,16 @@ func (api *API) updateEntitlements(ctx context.Context) error { if api.Options.DeploymentValues.DERP.Server.Enable { api.derpMesh.SetAddresses([]string{}, false) } - api.replicaManager.SetCallback(func() { + api.replicaManager.SetCallback("derp", func() { // If the amount of replicas change, so should our entitlements. // This is to display a warning in the UI if the user is unlicensed. _ = api.updateEntitlements(api.ctx) }) + + if natsPubsub, ok := api.Pubsub.(*nats.Pubsub); ok { + natsPubsub.SetPeerFetcher(nats.NopPeerFetcher{}) + api.replicaManager.SetCallback("nats", nil) + } } // Recheck changed in case the HA coordinator failed to set up. diff --git a/enterprise/coderd/coderd_test.go b/enterprise/coderd/coderd_test.go index 805b809699..7cdda8e64d 100644 --- a/enterprise/coderd/coderd_test.go +++ b/enterprise/coderd/coderd_test.go @@ -18,6 +18,7 @@ import ( "time" "github.com/google/uuid" + natsserver "github.com/nats-io/nats-server/v2/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/goleak" @@ -36,6 +37,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/entitlements" "github.com/coder/coder/v2/coderd/httpapi" agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds" @@ -43,6 +45,7 @@ import ( "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/util/namesgenerator" "github.com/coder/coder/v2/coderd/util/ptr" + natspubsub "github.com/coder/coder/v2/coderd/x/nats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/enterprise/audit" @@ -624,6 +627,95 @@ func TestMultiReplica_EmptyRelayAddress_DisabledDERP(t *testing.T) { } } +func TestMultiReplica_NATSPubsubPeers(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + db, pgPubsub := dbtestutil.NewDB(t) + clusterToken := "shared-token" + + natsA, err := natspubsub.New(ctx, logger.Named("nats-a"), natspubsub.Options{ + ClusterHost: "127.0.0.1", + ClusterPort: natsserver.RANDOM_PORT, + ClusterAuthToken: clusterToken, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = natsA.Close() }) + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentNATSPubsub)} + _, _ = coderdenttest.New(t, &coderdenttest.Options{ + EntitlementsUpdateInterval: 25 * time.Millisecond, + ReplicaSyncUpdateInterval: 25 * time.Millisecond, + Options: &coderdtest.Options{ + Logger: &logger, + Database: db, + Pubsub: natsA, + ReplicaSyncPubsub: pgPubsub.(*pubsub.PGPubsub), + DeploymentValues: dv, + }, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureHighAvailability: 1, + }, + }, + }) + + natsB, err := natspubsub.New(ctx, logger.Named("nats-b"), natspubsub.Options{ + ClusterHost: "127.0.0.1", + ClusterPort: natsserver.RANDOM_PORT, + ClusterAuthToken: clusterToken, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = natsB.Close() }) + + mgr, err := replicasync.New(ctx, logger.Named("replica-b"), db, pgPubsub, &replicasync.Options{ + ID: uuid.New(), + RelayAddress: fmt.Sprintf("nats://127.0.0.1:%d", natsB.Server.ClusterAddr().Port), + RegionID: 12345, + UpdateInterval: testutil.IntervalFast, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = mgr.Close() }) + + subject := "nats.replica" + messages := make(chan []byte, 1) + cancel, err := natsB.Subscribe(subject, func(_ context.Context, msg []byte) { + messages <- msg + }) + require.NoError(t, err) + defer cancel() + + payload := []byte("from-replicasync-peers") + var publishErr error + var flushErr error + var updateErr error + require.Eventually(t, func() bool { + updateErr = mgr.PublishUpdate() + if updateErr != nil { + return false + } + publishErr = natsA.Publish(subject, payload) + if publishErr != nil { + return false + } + flushErr = natsA.Flush() + if flushErr != nil { + return false + } + select { + case got := <-messages: + return string(got) == string(payload) + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast) + require.NoError(t, updateErr) + require.NoError(t, publishErr) + require.NoError(t, flushErr) +} + func TestSCIMDisabled(t *testing.T) { t.Parallel() diff --git a/enterprise/replicasync/replicasync.go b/enterprise/replicasync/replicasync.go index f69db6ed94..e7c067fff8 100644 --- a/enterprise/replicasync/replicasync.go +++ b/enterprise/replicasync/replicasync.go @@ -122,10 +122,10 @@ type Manager struct { closed chan (struct{}) closeCancel context.CancelFunc - self database.Replica - mutex sync.Mutex - peers []database.Replica - callback func() + self database.Replica + mutex sync.Mutex + peers []database.Replica + callbacks map[string]func() } func (m *Manager) ID() uuid.UUID { @@ -359,8 +359,8 @@ func (m *Manager) syncReplicas(ctx context.Context) error { } } m.self = replica - if m.callback != nil { - go m.callback() + for _, callback := range m.callbacks { + go callback() } return nil } @@ -414,6 +414,14 @@ func (m *Manager) AllPrimary() []database.Replica { return replicas } +func (m *Manager) PrimaryPeerAddresses() []string { + addresses := make([]string, 0, len(m.AllPrimary())) + for _, replica := range m.AllPrimary() { + addresses = append(addresses, replica.RelayAddress) + } + return addresses +} + // InRegion returns every replica in the given DERP region excluding itself. func (m *Manager) InRegion(regionID int32) []database.Replica { m.mutex.Lock() @@ -439,12 +447,20 @@ func (m *Manager) regionID() int32 { return m.self.RegionID } -// SetCallback sets a function to execute whenever new peers -// are refreshed or updated. -func (m *Manager) SetCallback(callback func()) { +// SetCallback sets a named function to execute whenever new peers are refreshed +// or updated. Calling SetCallback again with the same name replaces the prior +// callback. Passing nil removes the named callback. +func (m *Manager) SetCallback(name string, callback func()) { m.mutex.Lock() defer m.mutex.Unlock() - m.callback = callback + if callback == nil { + delete(m.callbacks, name) + return + } + if m.callbacks == nil { + m.callbacks = make(map[string]func()) + } + m.callbacks[name] = callback // Instantly call the callback to inform replicas! go callback() } diff --git a/enterprise/replicasync/replicasync_test.go b/enterprise/replicasync/replicasync_test.go index 0438db8e21..dfbd2fa2b1 100644 --- a/enterprise/replicasync/replicasync_test.go +++ b/enterprise/replicasync/replicasync_test.go @@ -207,6 +207,119 @@ func TestReplica(t *testing.T) { return len(server.Regional()) == 0 }, testutil.WaitShort, testutil.IntervalFast) }) + t.Run("MultipleCallbacks", func(t *testing.T) { + t.Parallel() + dh := &derpyHandler{} + defer dh.requireOnlyDERPPaths(t) + srv := httptest.NewServer(dh) + defer srv.Close() + db, pubsub := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + server, err := replicasync.New(ctx, testutil.Logger(t), db, pubsub, &replicasync.Options{ + RelayAddress: srv.URL, + }) + require.NoError(t, err) + defer server.Close() + + first := make(chan struct{}, 2) + second := make(chan struct{}, 2) + server.SetCallback("first", func() { first <- struct{}{} }) + server.SetCallback("second", func() { second <- struct{}{} }) + testutil.RequireReceive(ctx, t, first) + testutil.RequireReceive(ctx, t, second) + + require.NoError(t, server.UpdateNow(ctx)) + testutil.RequireReceive(ctx, t, first) + testutil.RequireReceive(ctx, t, second) + }) + t.Run("SetCallbackReplaces", func(t *testing.T) { + t.Parallel() + dh := &derpyHandler{} + defer dh.requireOnlyDERPPaths(t) + srv := httptest.NewServer(dh) + defer srv.Close() + db, pubsub := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + server, err := replicasync.New(ctx, testutil.Logger(t), db, pubsub, &replicasync.Options{ + RelayAddress: srv.URL, + }) + require.NoError(t, err) + defer server.Close() + + first := make(chan struct{}, 2) + second := make(chan struct{}, 2) + server.SetCallback("same", func() { first <- struct{}{} }) + testutil.RequireReceive(ctx, t, first) + + server.SetCallback("same", func() { second <- struct{}{} }) + testutil.RequireReceive(ctx, t, second) + require.NoError(t, server.UpdateNow(ctx)) + testutil.RequireReceive(ctx, t, second) + requireNoCallback(t, first) + }) + t.Run("SetCallbackDeletes", func(t *testing.T) { + t.Parallel() + dh := &derpyHandler{} + defer dh.requireOnlyDERPPaths(t) + srv := httptest.NewServer(dh) + defer srv.Close() + db, pubsub := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + server, err := replicasync.New(ctx, testutil.Logger(t), db, pubsub, &replicasync.Options{ + RelayAddress: srv.URL, + }) + require.NoError(t, err) + defer server.Close() + + called := make(chan struct{}, 2) + server.SetCallback("same", func() { called <- struct{}{} }) + testutil.RequireReceive(ctx, t, called) + + server.SetCallback("same", nil) + require.NoError(t, server.UpdateNow(ctx)) + requireNoCallback(t, called) + }) + t.Run("PrimaryPeerAddresses", func(t *testing.T) { + t.Parallel() + db, pubsub := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + primary, err := db.InsertReplica(ctx, database.InsertReplicaParams{ + ID: uuid.New(), + CreatedAt: dbtime.Now(), + StartedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + RelayAddress: "nats://primary.example:6222", + Primary: true, + }) + require.NoError(t, err) + _, err = db.InsertReplica(ctx, database.InsertReplicaParams{ + ID: uuid.New(), + CreatedAt: dbtime.Now(), + StartedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + RelayAddress: "nats://proxy.example:6222", + Primary: false, + }) + require.NoError(t, err) + _, err = db.InsertReplica(ctx, database.InsertReplicaParams{ + ID: uuid.New(), + CreatedAt: dbtime.Now(), + StartedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + Primary: true, + }) + require.NoError(t, err) + server, err := replicasync.New(ctx, testutil.Logger(t), db, pubsub, &replicasync.Options{ + RelayAddress: "nats://self.example:6222", + }) + require.NoError(t, err) + defer server.Close() + require.Contains(t, server.PrimaryPeerAddresses(), primary.RelayAddress) + require.ElementsMatch(t, []string{ + "nats://primary.example:6222", + "nats://self.example:6222", + }, server.PrimaryPeerAddresses()) + }) t.Run("TwentyConcurrent", func(t *testing.T) { // Ensures that twenty concurrent replicas can spawn and all // discover each other in parallel! @@ -233,7 +346,7 @@ func TestReplica(t *testing.T) { done := false var m sync.Mutex - server.SetCallback(func() { + server.SetCallback("all-primary", func() { m.Lock() defer m.Unlock() if len(server.AllPrimary()) != count { @@ -269,6 +382,15 @@ func TestReplica(t *testing.T) { }) } +func requireNoCallback(t *testing.T, ch <-chan struct{}) { + t.Helper() + select { + case <-ch: + require.FailNow(t, "unexpected callback") + default: + } +} + type derpyHandler struct { atomic.Uint32 } diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 5c0cf7c24d..807c63f400 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4392,6 +4392,7 @@ export type Experiment = | "auto-fill-parameters" | "example" | "mcp-server-http" + | "nats_pubsub" | "notifications" | "oauth2" | "workspace-build-updates" @@ -4401,6 +4402,7 @@ export const Experiments: Experiment[] = [ "auto-fill-parameters", "example", "mcp-server-http", + "nats_pubsub", "notifications", "oauth2", "workspace-build-updates", diff --git a/testutil/logger.go b/testutil/logger.go index 26cbde5655..4f3ca55d1d 100644 --- a/testutil/logger.go +++ b/testutil/logger.go @@ -32,6 +32,7 @@ func IgnoreLoggedError(entry slog.SinkEntry) bool { if xerrors.Is(err, yamux.ErrSessionShutdown) { return true } + // Canceled queries usually happen when we're shutting down tests, and so // ignoring them should reduce flakiness. This also includes // context.Canceled and context.DeadlineExceeded errors, even if they are