chore: refactor DERP setting loop (#15344)

Implements a Tailnet API DERP controller by refactoring from `workspacesdk`

chore re: #14729

tailnet/controllers.go

@@ -359,3 +359,84 @@ func NewInMemoryCoordinatorClient(
 	)
 	return c
 }
+
+type DERPMapSetter interface {
+	SetDERPMap(derpMap *tailcfg.DERPMap)
+}
+
+type basicDERPController struct {
+	logger slog.Logger
+	setter DERPMapSetter
+}
+
+func (b *basicDERPController) New(client DERPClient) CloserWaiter {
+	l := &derpSetLoop{
+		logger:       b.logger,
+		setter:       b.setter,
+		client:       client,
+		errChan:      make(chan error, 1),
+		recvLoopDone: make(chan struct{}),
+	}
+	go l.recvLoop()
+	return l
+}
+
+func NewBasicDERPController(logger slog.Logger, setter DERPMapSetter) DERPController {
+	return &basicDERPController{
+		logger: logger,
+		setter: setter,
+	}
+}
+
+type derpSetLoop struct {
+	logger slog.Logger
+	setter DERPMapSetter
+	client DERPClient
+
+	sync.Mutex
+	closed       bool
+	errChan      chan error
+	recvLoopDone chan struct{}
+}
+
+func (l *derpSetLoop) Close(ctx context.Context) error {
+	l.Lock()
+	defer l.Unlock()
+	if l.closed {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-l.recvLoopDone:
+			return nil
+		}
+	}
+	l.closed = true
+	cErr := l.client.Close()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-l.recvLoopDone:
+		return cErr
+	}
+}
+
+func (l *derpSetLoop) Wait() <-chan error {
+	return l.errChan
+}
+
+func (l *derpSetLoop) recvLoop() {
+	defer close(l.recvLoopDone)
+	for {
+		dm, err := l.client.Recv()
+		if err != nil {
+			l.logger.Debug(context.Background(), "failed to receive DERP message", slog.Error(err))
+			select {
+			case l.errChan <- err:
+			default:
+			}
+			return
+		}
+		l.logger.Debug(context.Background(), "got new DERP Map", slog.F("derp_map", dm))
+		l.setter.SetDERPMap(dm)
+	}
+}

tailnet/controllers_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
+	"golang.org/x/xerrors"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 
@@ -281,3 +282,72 @@ func (f *fakeCoordinatee) SetNodeCallback(callback func(*tailnet.Node)) {
 	defer f.Unlock()
 	f.callback = callback
 }
+
+func TestNewBasicDERPController_Mainline(t *testing.T) {
+	t.Parallel()
+	fs := make(chan *tailcfg.DERPMap)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	uut := tailnet.NewBasicDERPController(logger, fakeSetter(fs))
+	fc := fakeDERPClient{
+		ch: make(chan *tailcfg.DERPMap),
+	}
+	c := uut.New(fc)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	expectDM := &tailcfg.DERPMap{}
+	testutil.RequireSendCtx(ctx, t, fc.ch, expectDM)
+	gotDM := testutil.RequireRecvCtx(ctx, t, fs)
+	require.Equal(t, expectDM, gotDM)
+	err := c.Close(ctx)
+	require.NoError(t, err)
+	err = testutil.RequireRecvCtx(ctx, t, c.Wait())
+	require.ErrorIs(t, err, io.EOF)
+	// ensure Close is idempotent
+	err = c.Close(ctx)
+	require.NoError(t, err)
+}
+
+func TestNewBasicDERPController_RecvErr(t *testing.T) {
+	t.Parallel()
+	fs := make(chan *tailcfg.DERPMap)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	uut := tailnet.NewBasicDERPController(logger, fakeSetter(fs))
+	expectedErr := xerrors.New("a bad thing happened")
+	fc := fakeDERPClient{
+		ch:  make(chan *tailcfg.DERPMap),
+		err: expectedErr,
+	}
+	c := uut.New(fc)
+	ctx := testutil.Context(t, testutil.WaitShort)
+	err := testutil.RequireRecvCtx(ctx, t, c.Wait())
+	require.ErrorIs(t, err, expectedErr)
+	// ensure Close is idempotent
+	err = c.Close(ctx)
+	require.NoError(t, err)
+}
+
+type fakeSetter chan *tailcfg.DERPMap
+
+func (s fakeSetter) SetDERPMap(derpMap *tailcfg.DERPMap) {
+	s <- derpMap
+}
+
+type fakeDERPClient struct {
+	ch  chan *tailcfg.DERPMap
+	err error
+}
+
+func (f fakeDERPClient) Close() error {
+	close(f.ch)
+	return nil
+}
+
+func (f fakeDERPClient) Recv() (*tailcfg.DERPMap, error) {
+	if f.err != nil {
+		return nil, f.err
+	}
+	dm, ok := <-f.ch
+	if ok {
+		return dm, nil
+	}
+	return nil, io.EOF
+}

tailnet/convert.go

@@ -298,3 +298,21 @@ func WorkspaceStatusToProto(status codersdk.WorkspaceStatus) proto.Workspace_Sta
 		return proto.Workspace_UNKNOWN
 	}
 }
+
+type DERPFromDRPCWrapper struct {
+	Client proto.DRPCTailnet_StreamDERPMapsClient
+}
+
+func (w *DERPFromDRPCWrapper) Close() error {
+	return w.Client.Close()
+}
+
+func (w *DERPFromDRPCWrapper) Recv() (*tailcfg.DERPMap, error) {
+	p, err := w.Client.Recv()
+	if err != nil {
+		return nil, err
+	}
+	return DERPMapFromProto(p), nil
+}
+
+var _ DERPClient = &DERPFromDRPCWrapper{}