Skip to content

Commit fe63f6a

Browse files
spikecurtisspikecurtis
committed
feat: use tailnet v2 API for coordination
1 parent 04e3985 commit fe63f6a

30 files changed

+1187
-1108
lines changed

‎Makefile

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -475,7 +475,8 @@ gen: \
475475
site/.eslintignore \
476476
site/e2e/provisionerGenerated.ts \
477477
site/src/theme/icons.json \
478-
examples/examples.gen.json
478+
examples/examples.gen.json \
479+
tailnet/tailnettest/coordinatormock.go
479480
.PHONY: gen
480481

481482
# Mark all generated files as fresh so make thinks they're up-to-date. This is
@@ -502,6 +503,7 @@ gen/mark-fresh:
502503
site/e2e/provisionerGenerated.ts \
503504
site/src/theme/icons.json \
504505
examples/examples.gen.json \
506+
tailnet/tailnettest/coordinatormock.go \
505507
"
506508
for file in $$files; do
507509
echo "$$file"
@@ -529,6 +531,9 @@ coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $
529531
coderd/database/dbmock/dbmock.go: coderd/database/db.go coderd/database/querier.go
530532
go generate ./coderd/database/dbmock/
531533

534+
tailnet/tailnettest/coordinatormock.go: tailnet/coordinator.go
535+
go generate ./tailnet/tailnettest/
536+
532537
tailnet/proto/tailnet.pb.go: tailnet/proto/tailnet.proto
533538
protoc \
534539
--go_out=. \

‎agent/agent.go

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ import (
2222
"sync"
2323
"time"
2424

25+
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
26+
2527
"github.com/go-chi/chi/v5"
2628
"github.com/google/uuid"
2729
"github.com/prometheus/client_golang/prometheus"
@@ -30,6 +32,7 @@ import (
3032
"golang.org/x/exp/slices"
3133
"golang.org/x/sync/errgroup"
3234
"golang.org/x/xerrors"
35+
"storj.io/drpc"
3336
"tailscale.com/net/speedtest"
3437
"tailscale.com/tailcfg"
3538
"tailscale.com/types/netlogtype"
@@ -86,7 +89,7 @@ type Options struct {
8689

8790
type Client interface {
8891
Manifest(ctx context.Context) (agentsdk.Manifest, error)
89-
Listen(ctx context.Context) (net.Conn, error)
92+
Listen(ctx context.Context) (drpc.Conn, error)
9093
DERPMapUpdates(ctx context.Context) (<-chan agentsdk.DERPMapUpdate, io.Closer, error)
9194
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
9295
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
@@ -1058,20 +1061,34 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
10581061
ctx, cancel := context.WithCancel(ctx)
10591062
defer cancel()
10601063

1061-
coordinator, err := a.client.Listen(ctx)
1064+
conn, err := a.client.Listen(ctx)
10621065
if err != nil {
10631066
return err
10641067
}
1065-
defer coordinator.Close()
1068+
defer func() {
1069+
cErr := conn.Close()
1070+
if cErr != nil {
1071+
a.logger.Debug(ctx, "error closing drpc connection", slog.Error(err))
1072+
}
1073+
}()
1074+
1075+
tClient := tailnetproto.NewDRPCTailnetClient(conn)
1076+
coordinate, err := tClient.Coordinate(ctx)
1077+
if err != nil {
1078+
return xerrors.Errorf("failed to connect to the coordinate endpoint: %w", err)
1079+
}
1080+
defer func() {
1081+
cErr := coordinate.Close()
1082+
if cErr != nil {
1083+
a.logger.Debug(ctx, "error closing Coordinate client", slog.Error(err))
1084+
}
1085+
}()
10661086
a.logger.Info(ctx, "connected to coordination endpoint")
1067-
sendNodes, errChan := tailnet.ServeCoordinator(coordinator, func(nodes []*tailnet.Node) error {
1068-
return network.UpdateNodes(nodes, false)
1069-
})
1070-
network.SetNodeCallback(sendNodes)
1087+
coordination := tailnet.NewRemoteCoordination(a.logger, coordinate, network, uuid.Nil)
10711088
select {
10721089
case <-ctx.Done():
10731090
return ctx.Err()
1074-
case err := <-errChan:
1091+
case err := <-coordination.Error():
10751092
return err
10761093
}
10771094
}

‎agent/agent_test.go

Lines changed: 48 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1677,9 +1677,11 @@ func TestAgent_UpdatedDERP(t *testing.T) {
16771677
require.NotNil(t, originalDerpMap)
16781678

16791679
coordinator := tailnet.NewCoordinator(logger)
1680-
defer func() {
1680+
// use t.Cleanup so the coordinator closing doesn't deadlock with in-memory
1681+
// coordination
1682+
t.Cleanup(func() {
16811683
_ = coordinator.Close()
1682-
}()
1684+
})
16831685
agentID := uuid.New()
16841686
statsCh := make(chan *agentsdk.Stats, 50)
16851687
fs := afero.NewMemMapFs()
@@ -1694,41 +1696,42 @@ func TestAgent_UpdatedDERP(t *testing.T) {
16941696
statsCh,
16951697
coordinator,
16961698
)
1697-
closer := agent.New(agent.Options{
1699+
uut := agent.New(agent.Options{
16981700
Client: client,
16991701
Filesystem: fs,
17001702
Logger: logger.Named("agent"),
17011703
ReconnectingPTYTimeout: time.Minute,
17021704
})
1703-
defer func() {
1704-
_ = closer.Close()
1705-
}()
1705+
t.Cleanup(func() {
1706+
_ = uut.Close()
1707+
})
17061708

17071709
// Setup a client connection.
1708-
newClientConn := func(derpMap *tailcfg.DERPMap) *codersdk.WorkspaceAgentConn {
1710+
newClientConn := func(derpMap *tailcfg.DERPMap, name string) *codersdk.WorkspaceAgentConn {
17091711
conn, err := tailnet.NewConn(&tailnet.Options{
17101712
Addresses: []netip.Prefix{netip.PrefixFrom(tailnet.IP(), 128)},
17111713
DERPMap: derpMap,
1712-
Logger: logger.Named("client"),
1714+
Logger: logger.Named(name),
17131715
})
17141716
require.NoError(t, err)
1715-
clientConn, serverConn := net.Pipe()
1716-
serveClientDone := make(chan struct{})
17171717
t.Cleanup(func() {
1718-
_ = clientConn.Close()
1719-
_ = serverConn.Close()
1718+
t.Logf("closing conn %s", name)
17201719
_ = conn.Close()
1721-
<-serveClientDone
17221720
})
1723-
go func() {
1724-
defer close(serveClientDone)
1725-
err := coordinator.ServeClient(serverConn, uuid.New(), agentID)
1726-
assert.NoError(t, err)
1727-
}()
1728-
sendNode, _ := tailnet.ServeCoordinator(clientConn, func(nodes []*tailnet.Node) error {
1729-
return conn.UpdateNodes(nodes, false)
1721+
testCtx, testCtxCancel := context.WithCancel(context.Background())
1722+
t.Cleanup(testCtxCancel)
1723+
clientID := uuid.New()
1724+
coordination := tailnet.NewInMemoryCoordination(
1725+
testCtx, logger,
1726+
clientID, agentID,
1727+
coordinator, conn)
1728+
t.Cleanup(func() {
1729+
t.Logf("closing coordination %s", name)
1730+
err := coordination.Close()
1731+
if err != nil {
1732+
t.Logf("error closing in-memory coordination: %s", err.Error())
1733+
}
17301734
})
1731-
conn.SetNodeCallback(sendNode)
17321735
// Force DERP.
17331736
conn.SetBlockEndpoints(true)
17341737

@@ -1737,6 +1740,7 @@ func TestAgent_UpdatedDERP(t *testing.T) {
17371740
CloseFunc: func() error { return codersdk.ErrSkipClose },
17381741
})
17391742
t.Cleanup(func() {
1743+
t.Logf("closing sdkConn %s", name)
17401744
_ = sdkConn.Close()
17411745
})
17421746
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -1747,7 +1751,7 @@ func TestAgent_UpdatedDERP(t *testing.T) {
17471751

17481752
return sdkConn
17491753
}
1750-
conn1 := newClientConn(originalDerpMap)
1754+
conn1 := newClientConn(originalDerpMap, "client1")
17511755

17521756
// Change the DERP map.
17531757
newDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
@@ -1766,27 +1770,34 @@ func TestAgent_UpdatedDERP(t *testing.T) {
17661770
DERPMap: newDerpMap,
17671771
})
17681772
require.NoError(t, err)
1773+
t.Logf("client Pushed DERPMap update")
17691774

17701775
require.Eventually(t, func() bool {
1771-
conn := closer.TailnetConn()
1776+
conn := uut.TailnetConn()
17721777
if conn == nil {
17731778
return false
17741779
}
17751780
regionIDs := conn.DERPMap().RegionIDs()
1776-
return len(regionIDs) == 1 && regionIDs[0] == 2 && conn.Node().PreferredDERP == 2
1781+
preferredDERP := conn.Node().PreferredDERP
1782+
t.Logf("agent Conn DERPMap with regionIDs %v, PreferredDERP %d", regionIDs, preferredDERP)
1783+
return len(regionIDs) == 1 && regionIDs[0] == 2 && preferredDERP == 2
17771784
}, testutil.WaitLong, testutil.IntervalFast)
1785+
t.Logf("agent got the new DERPMap")
17781786

17791787
// Connect from a second client and make sure it uses the new DERP map.
1780-
conn2 := newClientConn(newDerpMap)
1788+
conn2 := newClientConn(newDerpMap, "client2")
17811789
require.Equal(t, []int{2}, conn2.DERPMap().RegionIDs())
1790+
t.Log("conn2 got the new DERPMap")
17821791

17831792
// If the first client gets a DERP map update, it should be able to
17841793
// reconnect just fine.
17851794
conn1.SetDERPMap(newDerpMap)
17861795
require.Equal(t, []int{2}, conn1.DERPMap().RegionIDs())
1796+
t.Log("set the new DERPMap on conn1")
17871797
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
17881798
defer cancel()
17891799
require.True(t, conn1.AwaitReachable(ctx))
1800+
t.Log("conn1 reached agent with new DERP")
17901801
}
17911802

17921803
func TestAgent_Speedtest(t *testing.T) {
@@ -2063,22 +2074,22 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
20632074
Logger: logger.Named("client"),
20642075
})
20652076
require.NoError(t, err)
2066-
clientConn, serverConn := net.Pipe()
2067-
serveClientDone := make(chan struct{})
20682077
t.Cleanup(func() {
2069-
_ = clientConn.Close()
2070-
_ = serverConn.Close()
20712078
_ = conn.Close()
2072-
<-serveClientDone
20732079
})
2074-
go func() {
2075-
defer close(serveClientDone)
2076-
coordinator.ServeClient(serverConn, uuid.New(), metadata.AgentID)
2077-
}()
2078-
sendNode, _ := tailnet.ServeCoordinator(clientConn, func(nodes []*tailnet.Node) error {
2079-
return conn.UpdateNodes(nodes, false)
2080+
testCtx, testCtxCancel := context.WithCancel(context.Background())
2081+
t.Cleanup(testCtxCancel)
2082+
clientID := uuid.New()
2083+
coordination := tailnet.NewInMemoryCoordination(
2084+
testCtx, logger,
2085+
clientID, metadata.AgentID,
2086+
coordinator, conn)
2087+
t.Cleanup(func() {
2088+
err := coordination.Close()
2089+
if err != nil {
2090+
t.Logf("error closing in-mem coordination: %s", err.Error())
2091+
}
20802092
})
2081-
conn.SetNodeCallback(sendNode)
20822093
agentConn := codersdk.NewWorkspaceAgentConn(conn, codersdk.WorkspaceAgentConnOptions{
20832094
AgentID: metadata.AgentID,
20842095
})

‎agent/agenttest/client.go

Lines changed: 50 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,28 @@ package agenttest
33
import (
44
"context"
55
"io"
6-
"net"
76
"sync"
7+
"sync/atomic"
88
"testing"
99
"time"
1010

11+
"github.com/stretchr/testify/require"
12+
13+
"storj.io/drpc"
14+
"storj.io/drpc/drpcmux"
15+
"storj.io/drpc/drpcserver"
16+
"tailscale.com/tailcfg"
17+
18+
"github.com/coder/coder/v2/tailnet/proto"
19+
1120
"github.com/google/uuid"
1221
"golang.org/x/exp/maps"
1322
"golang.org/x/xerrors"
1423

1524
"cdr.dev/slog"
1625
"github.com/coder/coder/v2/codersdk"
1726
"github.com/coder/coder/v2/codersdk/agentsdk"
27+
drpcsdk "github.com/coder/coder/v2/codersdk/drpc"
1828
"github.com/coder/coder/v2/tailnet"
1929
"github.com/coder/coder/v2/testutil"
2030
)
@@ -24,18 +34,39 @@ func NewClient(t testing.TB,
2434
agentID uuid.UUID,
2535
manifest agentsdk.Manifest,
2636
statsChan chan *agentsdk.Stats,
27-
coordinator tailnet.CoordinatorV1,
37+
coordinator tailnet.Coordinator,
2838
) *Client {
2939
if manifest.AgentID == uuid.Nil {
3040
manifest.AgentID = agentID
3141
}
42+
coordPtr := atomic.Pointer[tailnet.Coordinator]{}
43+
coordPtr.Store(&coordinator)
44+
mux := drpcmux.New()
45+
drpcService := &tailnet.DRPCService{
46+
CoordPtr: &coordPtr,
47+
Logger: logger,
48+
// TODO: handle DERPMap too!
49+
DerpMapUpdateFrequency: time.Hour,
50+
DerpMapFn: func() *tailcfg.DERPMap { panic("not implemented") },
51+
}
52+
err := proto.DRPCRegisterTailnet(mux, drpcService)
53+
require.NoError(t, err)
54+
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
55+
Log: func(err error) {
56+
if xerrors.Is(err, io.EOF) {
57+
return
58+
}
59+
logger.Debug(context.Background(), "drpc server error", slog.Error(err))
60+
},
61+
})
3262
return &Client{
3363
t: t,
3464
logger: logger.Named("client"),
3565
agentID: agentID,
3666
manifest: manifest,
3767
statsChan: statsChan,
3868
coordinator: coordinator,
69+
server: server,
3970
derpMapUpdates: make(chan agentsdk.DERPMapUpdate),
4071
}
4172
}
@@ -47,7 +78,8 @@ type Client struct {
4778
manifest agentsdk.Manifest
4879
metadata map[string]agentsdk.Metadata
4980
statsChan chan *agentsdk.Stats
50-
coordinator tailnet.CoordinatorV1
81+
coordinator tailnet.Coordinator
82+
server *drpcserver.Server
5183
LastWorkspaceAgent func()
5284
PatchWorkspaceLogs func() error
5385
GetServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
@@ -63,20 +95,29 @@ func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
6395
return c.manifest, nil
6496
}
6597

66-
func (c *Client) Listen(_ context.Context) (net.Conn, error) {
67-
clientConn, serverConn := net.Pipe()
98+
func (c *Client) Listen(_ context.Context) (drpc.Conn, error) {
99+
conn, lis := drpcsdk.MemTransportPipe()
68100
closed := make(chan struct{})
69101
c.LastWorkspaceAgent = func() {
70-
_ = serverConn.Close()
71-
_ = clientConn.Close()
102+
_ = conn.Close()
103+
_ = lis.Close()
72104
<-closed
73105
}
74106
c.t.Cleanup(c.LastWorkspaceAgent)
107+
serveCtx, cancel := context.WithCancel(context.Background())
108+
c.t.Cleanup(cancel)
109+
auth := tailnet.AgentTunnelAuth{}
110+
streamID := tailnet.StreamID{
111+
Name: "agenttest",
112+
ID: c.agentID,
113+
Auth: auth,
114+
}
115+
serveCtx = tailnet.WithStreamID(serveCtx, streamID)
75116
go func() {
76-
_ = c.coordinator.ServeAgent(serverConn, c.agentID, "")
117+
_ = c.server.Serve(serveCtx, lis)
77118
close(closed)
78119
}()
79-
return clientConn, nil
120+
return conn, nil
80121
}
81122

82123
func (c *Client) ReportStats(ctx context.Context, _ slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error) {

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy