mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: agent uses Tailnet v2 API for DERPMap updates (#11698)
Switches the Agent to use Tailnet v2 API to get DERPMap updates. Subsequent PRs will do the same for the CLI (`codersdk`) and `wsproxy`.
This commit is contained in:
+35
-38
@@ -89,7 +89,6 @@ type Options struct {
|
||||
type Client interface {
|
||||
Manifest(ctx context.Context) (agentsdk.Manifest, error)
|
||||
Listen(ctx context.Context) (drpc.Conn, error)
|
||||
DERPMapUpdates(ctx context.Context) (<-chan agentsdk.DERPMapUpdate, io.Closer, error)
|
||||
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
|
||||
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
|
||||
PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
|
||||
@@ -822,10 +821,22 @@ func (a *agent) run(ctx context.Context) error {
|
||||
network.SetBlockEndpoints(manifest.DisableDirectConnections)
|
||||
}
|
||||
|
||||
// Listen returns the dRPC connection we use for both Coordinator and DERPMap updates
|
||||
conn, err := a.client.Listen(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
cErr := conn.Close()
|
||||
if cErr != nil {
|
||||
a.logger.Debug(ctx, "error closing drpc connection", slog.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
eg, egCtx := errgroup.WithContext(ctx)
|
||||
eg.Go(func() error {
|
||||
a.logger.Debug(egCtx, "running tailnet connection coordinator")
|
||||
err := a.runCoordinator(egCtx, network)
|
||||
err := a.runCoordinator(egCtx, conn, network)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("run coordinator: %w", err)
|
||||
}
|
||||
@@ -834,7 +845,7 @@ func (a *agent) run(ctx context.Context) error {
|
||||
|
||||
eg.Go(func() error {
|
||||
a.logger.Debug(egCtx, "running derp map subscriber")
|
||||
err := a.runDERPMapSubscriber(egCtx, network)
|
||||
err := a.runDERPMapSubscriber(egCtx, conn, network)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("run derp map subscriber: %w", err)
|
||||
}
|
||||
@@ -1056,21 +1067,8 @@ func (a *agent) createTailnet(ctx context.Context, agentID uuid.UUID, derpMap *t
|
||||
|
||||
// runCoordinator runs a coordinator and returns whether a reconnect
|
||||
// should occur.
|
||||
func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
conn, err := a.client.Listen(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
cErr := conn.Close()
|
||||
if cErr != nil {
|
||||
a.logger.Debug(ctx, "error closing drpc connection", slog.Error(err))
|
||||
}
|
||||
}()
|
||||
|
||||
func (a *agent) runCoordinator(ctx context.Context, conn drpc.Conn, network *tailnet.Conn) error {
|
||||
defer a.logger.Debug(ctx, "disconnected from coordination RPC")
|
||||
tClient := tailnetproto.NewDRPCTailnetClient(conn)
|
||||
coordinate, err := tClient.Coordinate(ctx)
|
||||
if err != nil {
|
||||
@@ -1082,7 +1080,7 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
|
||||
a.logger.Debug(ctx, "error closing Coordinate client", slog.Error(err))
|
||||
}
|
||||
}()
|
||||
a.logger.Info(ctx, "connected to coordination endpoint")
|
||||
a.logger.Info(ctx, "connected to coordination RPC")
|
||||
coordination := tailnet.NewRemoteCoordination(a.logger, coordinate, network, uuid.Nil)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -1093,30 +1091,29 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
|
||||
}
|
||||
|
||||
// runDERPMapSubscriber runs a coordinator and returns if a reconnect should occur.
|
||||
func (a *agent) runDERPMapSubscriber(ctx context.Context, network *tailnet.Conn) error {
|
||||
func (a *agent) runDERPMapSubscriber(ctx context.Context, conn drpc.Conn, network *tailnet.Conn) error {
|
||||
defer a.logger.Debug(ctx, "disconnected from derp map RPC")
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
updates, closer, err := a.client.DERPMapUpdates(ctx)
|
||||
tClient := tailnetproto.NewDRPCTailnetClient(conn)
|
||||
stream, err := tClient.StreamDERPMaps(ctx, &tailnetproto.StreamDERPMapsRequest{})
|
||||
if err != nil {
|
||||
return err
|
||||
return xerrors.Errorf("stream DERP Maps: %w", err)
|
||||
}
|
||||
defer closer.Close()
|
||||
|
||||
a.logger.Info(ctx, "connected to derp map endpoint")
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case update := <-updates:
|
||||
if update.Err != nil {
|
||||
return update.Err
|
||||
}
|
||||
if update.DERPMap != nil && !tailnet.CompareDERPMaps(network.DERPMap(), update.DERPMap) {
|
||||
a.logger.Info(ctx, "updating derp map due to detected changes")
|
||||
network.SetDERPMap(update.DERPMap)
|
||||
}
|
||||
defer func() {
|
||||
cErr := stream.Close()
|
||||
if cErr != nil {
|
||||
a.logger.Debug(ctx, "error closing DERPMap stream", slog.Error(err))
|
||||
}
|
||||
}()
|
||||
a.logger.Info(ctx, "connected to derp map RPC")
|
||||
for {
|
||||
dmp, err := stream.Recv()
|
||||
if err != nil {
|
||||
return xerrors.Errorf("recv DERPMap error: %w", err)
|
||||
}
|
||||
dm := tailnet.DERPMapFromProto(dmp)
|
||||
network.SetDERPMap(dm)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+12
-4
@@ -1349,6 +1349,7 @@ func TestAgent_Lifecycle(t *testing.T) {
|
||||
make(chan *agentsdk.Stats, 50),
|
||||
tailnet.NewCoordinator(logger),
|
||||
)
|
||||
defer client.Close()
|
||||
|
||||
fs := afero.NewMemMapFs()
|
||||
agent := agent.New(agent.Options{
|
||||
@@ -1683,6 +1684,10 @@ func TestAgent_UpdatedDERP(t *testing.T) {
|
||||
statsCh,
|
||||
coordinator,
|
||||
)
|
||||
t.Cleanup(func() {
|
||||
t.Log("closing client")
|
||||
client.Close()
|
||||
})
|
||||
uut := agent.New(agent.Options{
|
||||
Client: client,
|
||||
Filesystem: fs,
|
||||
@@ -1690,6 +1695,7 @@ func TestAgent_UpdatedDERP(t *testing.T) {
|
||||
ReconnectingPTYTimeout: time.Minute,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
t.Log("closing agent")
|
||||
_ = uut.Close()
|
||||
})
|
||||
|
||||
@@ -1718,6 +1724,7 @@ func TestAgent_UpdatedDERP(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Logf("error closing in-memory coordination: %s", err.Error())
|
||||
}
|
||||
t.Logf("closed coordination %s", name)
|
||||
})
|
||||
// Force DERP.
|
||||
conn.SetBlockEndpoints(true)
|
||||
@@ -1753,11 +1760,9 @@ func TestAgent_UpdatedDERP(t *testing.T) {
|
||||
}
|
||||
|
||||
// Push a new DERP map to the agent.
|
||||
err := client.PushDERPMapUpdate(agentsdk.DERPMapUpdate{
|
||||
DERPMap: newDerpMap,
|
||||
})
|
||||
err := client.PushDERPMapUpdate(newDerpMap)
|
||||
require.NoError(t, err)
|
||||
t.Logf("client Pushed DERPMap update")
|
||||
t.Logf("pushed DERPMap update to agent")
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
conn := uut.TailnetConn()
|
||||
@@ -1826,6 +1831,7 @@ func TestAgent_Reconnect(t *testing.T) {
|
||||
statsCh,
|
||||
coordinator,
|
||||
)
|
||||
defer client.Close()
|
||||
initialized := atomic.Int32{}
|
||||
closer := agent.New(agent.Options{
|
||||
ExchangeToken: func(ctx context.Context) (string, error) {
|
||||
@@ -1862,6 +1868,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
|
||||
make(chan *agentsdk.Stats, 50),
|
||||
coordinator,
|
||||
)
|
||||
defer client.Close()
|
||||
filesystem := afero.NewMemMapFs()
|
||||
closer := agent.New(agent.Options{
|
||||
ExchangeToken: func(ctx context.Context) (string, error) {
|
||||
@@ -2039,6 +2046,7 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
|
||||
statsCh := make(chan *agentsdk.Stats, 50)
|
||||
fs := afero.NewMemMapFs()
|
||||
c := agenttest.NewClient(t, logger.Named("agent"), metadata.AgentID, metadata, statsCh, coordinator)
|
||||
t.Cleanup(c.Close)
|
||||
|
||||
options := agent.Options{
|
||||
Client: c,
|
||||
|
||||
+15
-21
@@ -39,12 +39,12 @@ func NewClient(t testing.TB,
|
||||
coordPtr := atomic.Pointer[tailnet.Coordinator]{}
|
||||
coordPtr.Store(&coordinator)
|
||||
mux := drpcmux.New()
|
||||
derpMapUpdates := make(chan *tailcfg.DERPMap)
|
||||
drpcService := &tailnet.DRPCService{
|
||||
CoordPtr: &coordPtr,
|
||||
Logger: logger,
|
||||
// TODO: handle DERPMap too!
|
||||
DerpMapUpdateFrequency: time.Hour,
|
||||
DerpMapFn: func() *tailcfg.DERPMap { panic("not implemented") },
|
||||
CoordPtr: &coordPtr,
|
||||
Logger: logger,
|
||||
DerpMapUpdateFrequency: time.Microsecond,
|
||||
DerpMapFn: func() *tailcfg.DERPMap { return <-derpMapUpdates },
|
||||
}
|
||||
err := proto.DRPCRegisterTailnet(mux, drpcService)
|
||||
require.NoError(t, err)
|
||||
@@ -64,7 +64,7 @@ func NewClient(t testing.TB,
|
||||
statsChan: statsChan,
|
||||
coordinator: coordinator,
|
||||
server: server,
|
||||
derpMapUpdates: make(chan agentsdk.DERPMapUpdate),
|
||||
derpMapUpdates: derpMapUpdates,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -85,23 +85,26 @@ type Client struct {
|
||||
lifecycleStates []codersdk.WorkspaceAgentLifecycle
|
||||
startup agentsdk.PostStartupRequest
|
||||
logs []agentsdk.Log
|
||||
derpMapUpdates chan agentsdk.DERPMapUpdate
|
||||
derpMapUpdates chan *tailcfg.DERPMap
|
||||
derpMapOnce sync.Once
|
||||
}
|
||||
|
||||
func (c *Client) Close() {
|
||||
c.derpMapOnce.Do(func() { close(c.derpMapUpdates) })
|
||||
}
|
||||
|
||||
func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
|
||||
return c.manifest, nil
|
||||
}
|
||||
|
||||
func (c *Client) Listen(_ context.Context) (drpc.Conn, error) {
|
||||
func (c *Client) Listen(ctx context.Context) (drpc.Conn, error) {
|
||||
conn, lis := drpcsdk.MemTransportPipe()
|
||||
closed := make(chan struct{})
|
||||
c.LastWorkspaceAgent = func() {
|
||||
_ = conn.Close()
|
||||
_ = lis.Close()
|
||||
<-closed
|
||||
}
|
||||
c.t.Cleanup(c.LastWorkspaceAgent)
|
||||
serveCtx, cancel := context.WithCancel(context.Background())
|
||||
serveCtx, cancel := context.WithCancel(ctx)
|
||||
c.t.Cleanup(cancel)
|
||||
auth := tailnet.AgentTunnelAuth{}
|
||||
streamID := tailnet.StreamID{
|
||||
@@ -112,7 +115,6 @@ func (c *Client) Listen(_ context.Context) (drpc.Conn, error) {
|
||||
serveCtx = tailnet.WithStreamID(serveCtx, streamID)
|
||||
go func() {
|
||||
_ = c.server.Serve(serveCtx, lis)
|
||||
close(closed)
|
||||
}()
|
||||
return conn, nil
|
||||
}
|
||||
@@ -235,7 +237,7 @@ func (c *Client) GetServiceBanner(ctx context.Context) (codersdk.ServiceBannerCo
|
||||
return codersdk.ServiceBannerConfig{}, nil
|
||||
}
|
||||
|
||||
func (c *Client) PushDERPMapUpdate(update agentsdk.DERPMapUpdate) error {
|
||||
func (c *Client) PushDERPMapUpdate(update *tailcfg.DERPMap) error {
|
||||
timer := time.NewTimer(testutil.WaitShort)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
@@ -247,14 +249,6 @@ func (c *Client) PushDERPMapUpdate(update agentsdk.DERPMapUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) DERPMapUpdates(_ context.Context) (<-chan agentsdk.DERPMapUpdate, io.Closer, error) {
|
||||
closed := make(chan struct{})
|
||||
return c.derpMapUpdates, closeFunc(func() error {
|
||||
close(closed)
|
||||
return nil
|
||||
}), nil
|
||||
}
|
||||
|
||||
type closeFunc func() error
|
||||
|
||||
func (c closeFunc) Close() error {
|
||||
|
||||
Reference in New Issue
Block a user