diff --git a/tailnet/controllers.go b/tailnet/controllers.go index 2328e19640a4d..b7d4e246a4bee 100644 --- a/tailnet/controllers.go +++ b/tailnet/controllers.go @@ -897,6 +897,21 @@ type Workspace struct { agents map[uuid.UUID]*Agent } +func (w *Workspace) Clone() Workspace { + agents := make(map[uuid.UUID]*Agent, len(w.agents)) + for k, v := range w.agents { + clone := v.Clone() + agents[k] = &clone + } + return Workspace{ + ID: w.ID, + Name: w.Name, + Status: w.Status, + ownerUsername: w.ownerUsername, + agents: agents, + } +} + type DNSNameOptions struct { Suffix string } @@ -1049,6 +1064,7 @@ func (t *tunnelUpdater) recvLoop() { t.logger.Debug(context.Background(), "tunnel updater recvLoop started") defer t.logger.Debug(context.Background(), "tunnel updater recvLoop done") defer close(t.recvLoopDone) + updateKind := Snapshot for { update, err := t.client.Recv() if err != nil { @@ -1061,8 +1077,10 @@ func (t *tunnelUpdater) recvLoop() { } t.logger.Debug(context.Background(), "got workspace update", slog.F("workspace_update", update), + slog.F("update_kind", updateKind), ) - err = t.handleUpdate(update) + err = t.handleUpdate(update, updateKind) + updateKind = Diff if err != nil { t.logger.Critical(context.Background(), "failed to handle workspace Update", slog.Error(err)) cErr := t.client.Close() @@ -1083,14 +1101,23 @@ type WorkspaceUpdate struct { UpsertedAgents []*Agent DeletedWorkspaces []*Workspace DeletedAgents []*Agent + Kind UpdateKind } +type UpdateKind int + +const ( + Diff UpdateKind = iota + Snapshot +) + func (w *WorkspaceUpdate) Clone() WorkspaceUpdate { clone := WorkspaceUpdate{ UpsertedWorkspaces: make([]*Workspace, len(w.UpsertedWorkspaces)), UpsertedAgents: make([]*Agent, len(w.UpsertedAgents)), DeletedWorkspaces: make([]*Workspace, len(w.DeletedWorkspaces)), DeletedAgents: make([]*Agent, len(w.DeletedAgents)), + Kind: w.Kind, } for i, ws := range w.UpsertedWorkspaces { clone.UpsertedWorkspaces[i] = &Workspace{ @@ -1115,7 +1142,7 @@ func (w *WorkspaceUpdate) Clone() WorkspaceUpdate { return clone } -func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error { +func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate, updateKind UpdateKind) error { t.Lock() defer t.Unlock() @@ -1124,6 +1151,7 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error { UpsertedAgents: []*Agent{}, DeletedWorkspaces: []*Workspace{}, DeletedAgents: []*Agent{}, + Kind: updateKind, } for _, uw := range update.UpsertedWorkspaces { diff --git a/tailnet/controllers_test.go b/tailnet/controllers_test.go index 67834de462655..bb5b543378ebe 100644 --- a/tailnet/controllers_test.go +++ b/tailnet/controllers_test.go @@ -1611,6 +1611,7 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) { }, DeletedWorkspaces: []*tailnet.Workspace{}, DeletedAgents: []*tailnet.Agent{}, + Kind: tailnet.Snapshot, } // And the callback @@ -1626,6 +1627,9 @@ func TestTunnelAllWorkspaceUpdatesController_Initial(t *testing.T) { slices.SortFunc(recvState.UpsertedAgents, func(a, b *tailnet.Agent) int { return strings.Compare(a.Name, b.Name) }) + // tunnel is still open, so it's a diff + currentState.Kind = tailnet.Diff + require.Equal(t, currentState, recvState) } @@ -1692,14 +1696,17 @@ func TestTunnelAllWorkspaceUpdatesController_DeleteAgent(t *testing.T) { }, DeletedWorkspaces: []*tailnet.Workspace{}, DeletedAgents: []*tailnet.Agent{}, + Kind: tailnet.Snapshot, } cbUpdate := testutil.TryReceive(ctx, t, fUH.ch) require.Equal(t, initRecvUp, cbUpdate) - // Current state should match initial state, err := updateCtrl.CurrentState() require.NoError(t, err) + // tunnel is still open, so it's a diff + initRecvUp.Kind = tailnet.Diff + require.Equal(t, initRecvUp, state) // Send update that removes w1a1 and adds w1a2 diff --git a/vpn/tunnel.go b/vpn/tunnel.go index 63de203980d14..6c71aecaa0965 100644 --- a/vpn/tunnel.go +++ b/vpn/tunnel.go @@ -88,6 +88,7 @@ func NewTunnel( netLoopDone: make(chan struct{}), uSendCh: s.sendCh, agents: map[uuid.UUID]tailnet.Agent{}, + workspaces: map[uuid.UUID]tailnet.Workspace{}, clock: quartz.NewReal(), }, } @@ -347,7 +348,9 @@ type updater struct { uSendCh chan<- *TunnelMessage // agents contains the agents that are currently connected to the tunnel. agents map[uuid.UUID]tailnet.Agent - conn Conn + // workspaces contains the workspaces to which agents are currently connected via the tunnel. + workspaces map[uuid.UUID]tailnet.Workspace + conn Conn clock quartz.Clock } @@ -397,6 +400,11 @@ func (u *updater) sendUpdateResponse(req *request[*TunnelMessage, *ManagerMessag // createPeerUpdateLocked creates a PeerUpdate message from a workspace update, populating // the network status of the agents. func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUpdate { + // if the update is a snapshot, we need to process the full state + if update.Kind == tailnet.Snapshot { + processSnapshotUpdate(&update, u.agents, u.workspaces) + } + out := &PeerUpdate{ UpsertedWorkspaces: make([]*Workspace, len(update.UpsertedWorkspaces)), UpsertedAgents: make([]*Agent, len(update.UpsertedAgents)), @@ -404,7 +412,20 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp DeletedAgents: make([]*Agent, len(update.DeletedAgents)), } - u.saveUpdateLocked(update) + // save the workspace update to the tunnel's state, such that it can + // be used to populate automated peer updates. + for _, agent := range update.UpsertedAgents { + u.agents[agent.ID] = agent.Clone() + } + for _, agent := range update.DeletedAgents { + delete(u.agents, agent.ID) + } + for _, workspace := range update.UpsertedWorkspaces { + u.workspaces[workspace.ID] = workspace.Clone() + } + for _, workspace := range update.DeletedWorkspaces { + delete(u.workspaces, workspace.ID) + } for i, ws := range update.UpsertedWorkspaces { out.UpsertedWorkspaces[i] = &Workspace{ @@ -413,6 +434,7 @@ func (u *updater) createPeerUpdateLocked(update tailnet.WorkspaceUpdate) *PeerUp Status: Workspace_Status(ws.Status), } } + upsertedAgents := u.convertAgentsLocked(update.UpsertedAgents) out.UpsertedAgents = upsertedAgents for i, ws := range update.DeletedWorkspaces { @@ -472,17 +494,6 @@ func (u *updater) convertAgentsLocked(agents []*tailnet.Agent) []*Agent { return out } -// saveUpdateLocked saves the workspace update to the tunnel's state, such that it can -// be used to populate automated peer updates. -func (u *updater) saveUpdateLocked(update tailnet.WorkspaceUpdate) { - for _, agent := range update.UpsertedAgents { - u.agents[agent.ID] = agent.Clone() - } - for _, agent := range update.DeletedAgents { - delete(u.agents, agent.ID) - } -} - // setConn sets the `conn` and returns false if there's already a connection set. func (u *updater) setConn(conn Conn) bool { u.mu.Lock() @@ -552,6 +563,46 @@ func (u *updater) netStatusLoop() { } } +// processSnapshotUpdate handles the logic when a full state update is received. +// When the tunnel is live, we only receive diffs, but the first packet on any given +// reconnect to the tailnet API is a full state. +// Without this logic we weren't processing deletes for any workspaces or agents deleted +// while the client was disconnected while the computer was asleep. +func processSnapshotUpdate(update *tailnet.WorkspaceUpdate, agents map[uuid.UUID]tailnet.Agent, workspaces map[uuid.UUID]tailnet.Workspace) { + // ignoredWorkspaces is initially populated with the workspaces that are + // in the current update. Later on we populate it with the deleted workspaces too + // so that we don't send duplicate updates. Same applies to ignoredAgents. + ignoredWorkspaces := make(map[uuid.UUID]struct{}, len(update.UpsertedWorkspaces)) + ignoredAgents := make(map[uuid.UUID]struct{}, len(update.UpsertedAgents)) + + for _, workspace := range update.UpsertedWorkspaces { + ignoredWorkspaces[workspace.ID] = struct{}{} + } + for _, agent := range update.UpsertedAgents { + ignoredAgents[agent.ID] = struct{}{} + } + for _, agent := range agents { + if _, present := ignoredAgents[agent.ID]; !present { + // delete any current agents that are not in the new update + update.DeletedAgents = append(update.DeletedAgents, &tailnet.Agent{ + ID: agent.ID, + Name: agent.Name, + WorkspaceID: agent.WorkspaceID, + }) + } + } + for _, workspace := range workspaces { + if _, present := ignoredWorkspaces[workspace.ID]; !present { + update.DeletedWorkspaces = append(update.DeletedWorkspaces, &tailnet.Workspace{ + ID: workspace.ID, + Name: workspace.Name, + Status: workspace.Status, + }) + ignoredWorkspaces[workspace.ID] = struct{}{} + } + } +} + // hostsToIPStrings returns a slice of all unique IP addresses in the values // of the given map. func hostsToIPStrings(hosts map[dnsname.FQDN][]netip.Addr) []string { diff --git a/vpn/tunnel_internal_test.go b/vpn/tunnel_internal_test.go index d1d7377361f79..2beba66d7a7d2 100644 --- a/vpn/tunnel_internal_test.go +++ b/vpn/tunnel_internal_test.go @@ -2,6 +2,7 @@ package vpn import ( "context" + "maps" "net" "net/netip" "net/url" @@ -292,6 +293,7 @@ func TestUpdater_createPeerUpdate(t *testing.T) { ctx: ctx, netLoopDone: make(chan struct{}), agents: map[uuid.UUID]tailnet.Agent{}, + workspaces: map[uuid.UUID]tailnet.Workspace{}, conn: newFakeConn(tailnet.WorkspaceUpdate{}, hsTime), } @@ -486,6 +488,212 @@ func TestTunnel_sendAgentUpdate(t *testing.T) { require.Equal(t, hsTime, req.msg.GetPeerUpdate().UpsertedAgents[0].LastHandshake.AsTime()) } +func TestTunnel_sendAgentUpdateReconnect(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + + mClock := quartz.NewMock(t) + + wID1 := uuid.UUID{1} + aID1 := uuid.UUID{2} + aID2 := uuid.UUID{3} + + hsTime := time.Now().Add(-time.Minute).UTC() + + client := newFakeClient(ctx, t) + conn := newFakeConn(tailnet.WorkspaceUpdate{}, hsTime) + + tun, mgr := setupTunnel(t, ctx, client, mClock) + errCh := make(chan error, 1) + var resp *TunnelMessage + go func() { + r, err := mgr.unaryRPC(ctx, &ManagerMessage{ + Msg: &ManagerMessage_Start{ + Start: &StartRequest{ + TunnelFileDescriptor: 2, + CoderUrl: "https://coder.example.com", + ApiToken: "fakeToken", + }, + }, + }) + resp = r + errCh <- err + }() + testutil.RequireSend(ctx, t, client.ch, conn) + err := testutil.TryReceive(ctx, t, errCh) + require.NoError(t, err) + _, ok := resp.Msg.(*TunnelMessage_Start) + require.True(t, ok) + + // Inform the tunnel of the initial state + err = tun.Update(tailnet.WorkspaceUpdate{ + UpsertedWorkspaces: []*tailnet.Workspace{ + { + ID: wID1, Name: "w1", Status: proto.Workspace_STARTING, + }, + }, + UpsertedAgents: []*tailnet.Agent{ + { + ID: aID1, + Name: "agent1", + WorkspaceID: wID1, + Hosts: map[dnsname.FQDN][]netip.Addr{ + "agent1.coder.": {netip.MustParseAddr("fd60:627a:a42b:0101::")}, + }, + }, + }, + }) + require.NoError(t, err) + req := testutil.TryReceive(ctx, t, mgr.requests) + require.Nil(t, req.msg.Rpc) + require.NotNil(t, req.msg.GetPeerUpdate()) + require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 1) + require.Equal(t, aID1[:], req.msg.GetPeerUpdate().UpsertedAgents[0].Id) + + // Upsert a new agent simulating a reconnect + err = tun.Update(tailnet.WorkspaceUpdate{ + UpsertedWorkspaces: []*tailnet.Workspace{ + { + ID: wID1, Name: "w1", Status: proto.Workspace_STARTING, + }, + }, + UpsertedAgents: []*tailnet.Agent{ + { + ID: aID2, + Name: "agent2", + WorkspaceID: wID1, + Hosts: map[dnsname.FQDN][]netip.Addr{ + "agent2.coder.": {netip.MustParseAddr("fd60:627a:a42b:0101::")}, + }, + }, + }, + Kind: tailnet.Snapshot, + }) + require.NoError(t, err) + + // The new update only contains the new agent + mClock.AdvanceNext() + req = testutil.TryReceive(ctx, t, mgr.requests) + require.Nil(t, req.msg.Rpc) + peerUpdate := req.msg.GetPeerUpdate() + require.NotNil(t, peerUpdate) + require.Len(t, peerUpdate.UpsertedAgents, 1) + require.Len(t, peerUpdate.DeletedAgents, 1) + require.Len(t, peerUpdate.DeletedWorkspaces, 0) + + require.Equal(t, aID2[:], peerUpdate.UpsertedAgents[0].Id) + require.Equal(t, hsTime, peerUpdate.UpsertedAgents[0].LastHandshake.AsTime()) + + require.Equal(t, aID1[:], peerUpdate.DeletedAgents[0].Id) +} + +func TestTunnel_sendAgentUpdateWorkspaceReconnect(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + + mClock := quartz.NewMock(t) + + wID1 := uuid.UUID{1} + wID2 := uuid.UUID{2} + aID1 := uuid.UUID{3} + aID3 := uuid.UUID{4} + + hsTime := time.Now().Add(-time.Minute).UTC() + + client := newFakeClient(ctx, t) + conn := newFakeConn(tailnet.WorkspaceUpdate{}, hsTime) + + tun, mgr := setupTunnel(t, ctx, client, mClock) + errCh := make(chan error, 1) + var resp *TunnelMessage + go func() { + r, err := mgr.unaryRPC(ctx, &ManagerMessage{ + Msg: &ManagerMessage_Start{ + Start: &StartRequest{ + TunnelFileDescriptor: 2, + CoderUrl: "https://coder.example.com", + ApiToken: "fakeToken", + }, + }, + }) + resp = r + errCh <- err + }() + testutil.RequireSend(ctx, t, client.ch, conn) + err := testutil.TryReceive(ctx, t, errCh) + require.NoError(t, err) + _, ok := resp.Msg.(*TunnelMessage_Start) + require.True(t, ok) + + // Inform the tunnel of the initial state + err = tun.Update(tailnet.WorkspaceUpdate{ + UpsertedWorkspaces: []*tailnet.Workspace{ + { + ID: wID1, Name: "w1", Status: proto.Workspace_STARTING, + }, + }, + UpsertedAgents: []*tailnet.Agent{ + { + ID: aID1, + Name: "agent1", + WorkspaceID: wID1, + Hosts: map[dnsname.FQDN][]netip.Addr{ + "agent1.coder.": {netip.MustParseAddr("fd60:627a:a42b:0101::")}, + }, + }, + }, + }) + require.NoError(t, err) + req := testutil.TryReceive(ctx, t, mgr.requests) + require.Nil(t, req.msg.Rpc) + require.NotNil(t, req.msg.GetPeerUpdate()) + require.Len(t, req.msg.GetPeerUpdate().UpsertedAgents, 1) + require.Equal(t, aID1[:], req.msg.GetPeerUpdate().UpsertedAgents[0].Id) + + // Upsert a new agent with a new workspace while simulating a reconnect + err = tun.Update(tailnet.WorkspaceUpdate{ + UpsertedWorkspaces: []*tailnet.Workspace{ + { + ID: wID2, Name: "w2", Status: proto.Workspace_STARTING, + }, + }, + UpsertedAgents: []*tailnet.Agent{ + { + ID: aID3, + Name: "agent3", + WorkspaceID: wID2, + Hosts: map[dnsname.FQDN][]netip.Addr{ + "agent3.coder.": {netip.MustParseAddr("fd60:627a:a42b:0101::")}, + }, + }, + }, + Kind: tailnet.Snapshot, + }) + require.NoError(t, err) + + // The new update only contains the new agent + mClock.AdvanceNext() + req = testutil.TryReceive(ctx, t, mgr.requests) + mClock.AdvanceNext() + + require.Nil(t, req.msg.Rpc) + peerUpdate := req.msg.GetPeerUpdate() + require.NotNil(t, peerUpdate) + require.Len(t, peerUpdate.UpsertedWorkspaces, 1) + require.Len(t, peerUpdate.UpsertedAgents, 1) + require.Len(t, peerUpdate.DeletedAgents, 1) + require.Len(t, peerUpdate.DeletedWorkspaces, 1) + + require.Equal(t, wID2[:], peerUpdate.UpsertedWorkspaces[0].Id) + require.Equal(t, aID3[:], peerUpdate.UpsertedAgents[0].Id) + require.Equal(t, hsTime, peerUpdate.UpsertedAgents[0].LastHandshake.AsTime()) + + require.Equal(t, aID1[:], peerUpdate.DeletedAgents[0].Id) + require.Equal(t, wID1[:], peerUpdate.DeletedWorkspaces[0].Id) +} + //nolint:revive // t takes precedence func setupTunnel(t *testing.T, ctx context.Context, client *fakeClient, mClock quartz.Clock) (*Tunnel, *speaker[*ManagerMessage, *TunnelMessage, TunnelMessage]) { mp, tp := net.Pipe() @@ -513,3 +721,191 @@ func setupTunnel(t *testing.T, ctx context.Context, client *fakeClient, mClock q mgr.start() return tun, mgr } + +func TestProcessFreshState(t *testing.T) { + t.Parallel() + + wsID1 := uuid.New() + wsID2 := uuid.New() + wsID3 := uuid.New() + wsID4 := uuid.New() + + agentID1 := uuid.New() + agentID2 := uuid.New() + agentID3 := uuid.New() + agentID4 := uuid.New() + + agent1 := tailnet.Agent{ID: agentID1, Name: "agent1", WorkspaceID: wsID1} + agent2 := tailnet.Agent{ID: agentID2, Name: "agent2", WorkspaceID: wsID2} + agent3 := tailnet.Agent{ID: agentID3, Name: "agent3", WorkspaceID: wsID3} + agent4 := tailnet.Agent{ID: agentID4, Name: "agent4", WorkspaceID: wsID1} + + ws1 := tailnet.Workspace{ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING} + ws2 := tailnet.Workspace{ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING} + ws3 := tailnet.Workspace{ID: wsID3, Name: "ws3", Status: proto.Workspace_RUNNING} + ws4 := tailnet.Workspace{ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING} + + initialAgents := map[uuid.UUID]tailnet.Agent{ + agentID1: agent1, + agentID2: agent2, + agentID4: agent4, + } + initialWorkspaces := map[uuid.UUID]tailnet.Workspace{ + wsID1: ws1, + wsID2: ws2, + } + + tests := []struct { + name string + initialAgents map[uuid.UUID]tailnet.Agent + initialWorkspaces map[uuid.UUID]tailnet.Workspace + update *tailnet.WorkspaceUpdate + expectedDelete *tailnet.WorkspaceUpdate // We only care about deletions added by the function + }{ + { + name: "NoChange", + initialAgents: initialAgents, + initialWorkspaces: initialWorkspaces, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, + UpsertedAgents: []*tailnet.Agent{&agent1, &agent2, &agent4}, + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ // Expect no *additional* deletions + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + }, + { + name: "AgentAdded", // Agent 3 added in update + initialAgents: initialAgents, + initialWorkspaces: initialWorkspaces, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2, &ws3}, + UpsertedAgents: []*tailnet.Agent{&agent1, &agent2, &agent3, &agent4}, + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + }, + { + name: "AgentRemovedWorkspaceAlsoRemoved", // Agent 2 removed, ws2 also removed + initialAgents: initialAgents, + initialWorkspaces: initialWorkspaces, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1}, // ws2 not present + UpsertedAgents: []*tailnet.Agent{&agent1, &agent4}, // agent2 not present + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ + DeletedWorkspaces: []*tailnet.Workspace{ + {ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING}, + }, // Expect ws2 to be deleted + DeletedAgents: []*tailnet.Agent{ // Expect agent2 to be deleted + {ID: agentID2, Name: "agent2", WorkspaceID: wsID2}, + }, + }, + }, + { + name: "AgentRemovedWorkspaceStays", // Agent 4 removed, but ws1 stays (due to agent1) + initialAgents: initialAgents, + initialWorkspaces: initialWorkspaces, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, // ws1 still present + UpsertedAgents: []*tailnet.Agent{&agent1, &agent2}, // agent4 not present + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ + DeletedWorkspaces: []*tailnet.Workspace{}, // ws1 should NOT be deleted + DeletedAgents: []*tailnet.Agent{ // Expect agent4 to be deleted + {ID: agentID4, Name: "agent4", WorkspaceID: wsID1}, + }, + }, + }, + { + name: "InitialAgentsEmpty", + initialAgents: map[uuid.UUID]tailnet.Agent{}, // Start with no agents known + initialWorkspaces: map[uuid.UUID]tailnet.Workspace{}, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, + UpsertedAgents: []*tailnet.Agent{&agent1, &agent2}, + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ // Expect no deletions added + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + }, + { + name: "UpdateEmpty", // Snapshot says nothing exists + initialAgents: initialAgents, + initialWorkspaces: initialWorkspaces, + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{}, + UpsertedAgents: []*tailnet.Agent{}, + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted + DeletedWorkspaces: []*tailnet.Workspace{ + {ID: wsID1, Name: "ws1", Status: proto.Workspace_RUNNING}, + {ID: wsID2, Name: "ws2", Status: proto.Workspace_RUNNING}, + }, // ws1 and ws2 deleted + DeletedAgents: []*tailnet.Agent{ // agent1, agent2, agent4 deleted + {ID: agentID1, Name: "agent1", WorkspaceID: wsID1}, + {ID: agentID2, Name: "agent2", WorkspaceID: wsID2}, + {ID: agentID4, Name: "agent4", WorkspaceID: wsID1}, + }, + }, + }, + { + name: "WorkspaceWithNoAgents", // Snapshot says nothing exists + initialAgents: initialAgents, + initialWorkspaces: map[uuid.UUID]tailnet.Workspace{wsID1: ws1, wsID2: ws2, wsID4: ws4}, // ws4 has no agents + update: &tailnet.WorkspaceUpdate{ + Kind: tailnet.Snapshot, + UpsertedWorkspaces: []*tailnet.Workspace{&ws1, &ws2}, + UpsertedAgents: []*tailnet.Agent{&agent1, &agent2, &agent4}, + DeletedWorkspaces: []*tailnet.Workspace{}, + DeletedAgents: []*tailnet.Agent{}, + }, + expectedDelete: &tailnet.WorkspaceUpdate{ // Expect all initial agents/workspaces to be deleted + DeletedWorkspaces: []*tailnet.Workspace{ + {ID: wsID4, Name: "ws4", Status: proto.Workspace_RUNNING}, + }, // ws4 should be deleted + DeletedAgents: []*tailnet.Agent{}, + }, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + agentsCopy := make(map[uuid.UUID]tailnet.Agent) + maps.Copy(agentsCopy, tt.initialAgents) + + workspaceCopy := make(map[uuid.UUID]tailnet.Workspace) + maps.Copy(workspaceCopy, tt.initialWorkspaces) + + processSnapshotUpdate(tt.update, agentsCopy, workspaceCopy) + + require.ElementsMatch(t, tt.expectedDelete.DeletedAgents, tt.update.DeletedAgents, "DeletedAgents mismatch") + require.ElementsMatch(t, tt.expectedDelete.DeletedWorkspaces, tt.update.DeletedWorkspaces, "DeletedWorkspaces mismatch") + }) + } +}