fix: update conn derpmap every 5s in single tailnet (#9176)

deansheather · web-flow · commit f35423c041de · 2023-08-23T11:20:31.000Z
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -404,7 +404,7 @@ func New(options *Options) *API {
 		api.agentProvider, err = NewServerTailnet(api.ctx,
 			options.Logger,
 			options.DERPServer,
-			options.BaseDERPMap,
+			api.DERPMap,
 			func(context.Context) (tailnet.MultiAgentConn, error) {
 				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
 			},
diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -44,25 +44,49 @@ func NewServerTailnet(
 	ctx context.Context,
 	logger slog.Logger,
 	derpServer *derp.Server,
-	derpMap *tailcfg.DERPMap,
+	derpMapFn func() *tailcfg.DERPMap,
 	getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error),
 	cache *wsconncache.Cache,
 	traceProvider trace.TracerProvider,
 ) (*ServerTailnet, error) {
 	logger = logger.Named("servertailnet")
+	originalDerpMap := derpMapFn()
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
-		DERPMap:   derpMap,
+		DERPMap:   originalDerpMap,
 		Logger:    logger,
 	})
 	if err != nil {
 		return nil, xerrors.Errorf("create tailnet conn: %w", err)
 	}
 
 	serverCtx, cancel := context.WithCancel(ctx)
+	derpMapUpdaterClosed := make(chan struct{})
+	go func() {
+		defer close(derpMapUpdaterClosed)
+
+		ticker := time.NewTicker(5 * time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-serverCtx.Done():
+				return
+			case <-ticker.C:
+			}
+
+			newDerpMap := derpMapFn()
+			if !tailnet.CompareDERPMaps(originalDerpMap, newDerpMap) {
+				conn.SetDERPMap(newDerpMap)
+				originalDerpMap = newDerpMap
+			}
+		}
+	}()
+
 	tn := &ServerTailnet{
 		ctx:                  serverCtx,
 		cancel:               cancel,
+		derpMapUpdaterClosed: derpMapUpdaterClosed,
 		logger:               logger,
 		tracer:               traceProvider.Tracer(tracing.TracerName),
 		conn:                 conn,
@@ -237,8 +261,9 @@ func (s *ServerTailnet) reinitCoordinator() {
 }
 
 type ServerTailnet struct {
-	ctx    context.Context
-	cancel func()
+	ctx                  context.Context
+	cancel               func()
+	derpMapUpdaterClosed chan struct{}
 
 	logger        slog.Logger
 	tracer        trace.Tracer
@@ -403,5 +428,6 @@ func (s *ServerTailnet) Close() error {
 	_ = s.cache.Close()
 	_ = s.conn.Close()
 	s.transport.CloseIdleConnections()
+	<-s.derpMapUpdaterClosed
 	return nil
 }
diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -15,6 +15,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel/trace"
+	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
@@ -230,7 +231,7 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
 		context.Background(),
 		logger,
 		derpServer,
-		manifest.DERPMap,
+		func() *tailcfg.DERPMap { return manifest.DERPMap },
 		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
 		cache,
 		trace.NewNoopTracerProvider(),
diff --git a/docs/api/schemas.md b/docs/api/schemas.md
diff --git a/enterprise/coderd/workspaceproxy.go b/enterprise/coderd/workspaceproxy.go
@@ -720,6 +720,7 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 		AppSecurityKey:  api.AppSecurityKey.String(),
 		DERPMeshKey:     api.DERPServer.MeshKey(),
 		DERPRegionID:    regionID,
+		DERPMap:         api.AGPL.DERPMap(),
 		SiblingReplicas: siblingsRes,
 	})
 
diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -21,6 +21,7 @@ import (
 	"golang.org/x/xerrors"
 	"tailscale.com/derp"
 	"tailscale.com/derp/derphttp"
+	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 
 	"cdr.dev/slog"
@@ -119,7 +120,8 @@ type Server struct {
 	SDKClient *wsproxysdk.Client
 
 	// DERP
-	derpMesh *derpmesh.Mesh
+	derpMesh      *derpmesh.Mesh
+	latestDERPMap *tailcfg.DERPMap
 
 	// Used for graceful shutdown. Required for the dialer.
 	ctx           context.Context
@@ -239,17 +241,14 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 		return nil, xerrors.Errorf("parse app security key: %w", err)
 	}
 
-	connInfo, err := client.SDKClient.WorkspaceAgentConnectionInfoGeneric(ctx)
-	if err != nil {
-		return nil, xerrors.Errorf("get derpmap: %w", err)
-	}
-
 	var agentProvider workspaceapps.AgentProvider
 	if opts.Experiments.Enabled(codersdk.ExperimentSingleTailnet) {
 		stn, err := coderd.NewServerTailnet(ctx,
 			s.Logger,
 			nil,
-			connInfo.DERPMap,
+			func() *tailcfg.DERPMap {
+				return s.latestDERPMap
+			},
 			s.DialCoordinator,
 			wsconncache.New(s.DialWorkspaceAgent, 0),
 			s.TracerProvider,
@@ -456,6 +455,8 @@ func (s *Server) handleRegister(_ context.Context, res wsproxysdk.RegisterWorksp
 	}
 	s.derpMesh.SetAddresses(addresses, false)
 
+	s.latestDERPMap = res.DERPMap
+
 	return nil
 }
 
diff --git a/enterprise/wsproxy/wsproxysdk/wsproxysdk.go b/enterprise/wsproxy/wsproxysdk/wsproxysdk.go
@@ -14,6 +14,7 @@ import (
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
+	"tailscale.com/tailcfg"
 	"tailscale.com/util/singleflight"
 
 	"cdr.dev/slog"
@@ -206,9 +207,10 @@ type RegisterWorkspaceProxyRequest struct {
 }
 
 type RegisterWorkspaceProxyResponse struct {
-	AppSecurityKey string `json:"app_security_key"`
-	DERPMeshKey    string `json:"derp_mesh_key"`
-	DERPRegionID   int32  `json:"derp_region_id"`
+	AppSecurityKey string           `json:"app_security_key"`
+	DERPMeshKey    string           `json:"derp_mesh_key"`
+	DERPRegionID   int32            `json:"derp_region_id"`
+	DERPMap        *tailcfg.DERPMap `json:"derp_map"`
 	// SiblingReplicas is a list of all other replicas of the proxy that have
 	// not timed out.
 	SiblingReplicas []codersdk.Replica `json:"sibling_replicas"`
