Skip to content

Commit f7aad62

Browse files
ethanndicksonethanndickson
committed
review p2
1 parent 71a6e59 commit f7aad62

File tree

8 files changed

+333
-151
lines changed

8 files changed

+333
-151
lines changed

‎tailnet/controllers.go

Lines changed: 156 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ type DNSHostsSetter interface {
114114

115115
// UpdatesHandler is anything that expects a stream of workspace update diffs.
116116
type UpdatesHandler interface {
117-
Update(*proto.WorkspaceUpdate) error
117+
Update(WorkspaceUpdate) error
118118
}
119119

120120
// ControlProtocolClients represents an abstract interface to the tailnet control plane via a set
@@ -871,43 +871,51 @@ type TunnelAllWorkspaceUpdatesController struct {
871871
updater *tunnelUpdater
872872
}
873873

874-
type workspace struct {
875-
id uuid.UUID
876-
name string
877-
agents map[uuid.UUID]agent
874+
type Workspace struct {
875+
ID uuid.UUID
876+
Name string
877+
Status proto.Workspace_Status
878+
879+
ownerUsername string
880+
agents map[uuid.UUID]*Agent
878881
}
879882

880-
// addAllDNSNames adds names for all of its agents to the given map of names
881-
func (w workspace) addAllDNSNames(names map[dnsname.FQDN][]netip.Addr, owner string) error {
882-
for _, a := range w.agents {
883+
// updateDNSNames updates the DNS names for all agents in the workspace.
884+
func (w *Workspace) updateDNSNames() error {
885+
for id, a := range w.agents {
886+
names := make(map[dnsname.FQDN][]netip.Addr)
883887
// TODO: technically, DNS labels cannot start with numbers, but the rules are often not
884888
// strictly enforced.
885-
fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%s.%s.me.coder.", a.name, w.name))
886-
if err != nil {
887-
return err
888-
}
889-
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.id)}
890-
fqdn, err = dnsname.ToFQDN(fmt.Sprintf("%s.%s.%s.coder.", a.name, w.name, owner))
889+
fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%s.%s.me.coder.", a.Name, w.Name))
891890
if err != nil {
892891
return err
893892
}
894-
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.id)}
895-
}
896-
if len(w.agents) == 1 {
897-
fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%s.coder.", w.name))
893+
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.ID)}
894+
fqdn, err = dnsname.ToFQDN(fmt.Sprintf("%s.%s.%s.coder.", a.Name, w.Name, w.ownerUsername))
898895
if err != nil {
899896
return err
900897
}
901-
for _, a := range w.agents {
902-
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.id)}
898+
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.ID)}
899+
if len(w.agents) == 1 {
900+
fqdn, err := dnsname.ToFQDN(fmt.Sprintf("%s.coder.", w.Name))
901+
if err != nil {
902+
return err
903+
}
904+
for _, a := range w.agents {
905+
names[fqdn] = []netip.Addr{CoderServicePrefix.AddrFromUUID(a.ID)}
906+
}
903907
}
908+
a.Hosts = names
909+
w.agents[id] = a
904910
}
905911
return nil
906912
}
907913

908-
type agent struct {
909-
id uuid.UUID
910-
name string
914+
type Agent struct {
915+
ID uuid.UUID
916+
Name string
917+
WorkspaceID uuid.UUID
918+
Hosts map[dnsname.FQDN][]netip.Addr
911919
}
912920

913921
func (t *TunnelAllWorkspaceUpdatesController) New(client WorkspaceUpdatesClient) CloserWaiter {
@@ -922,40 +930,43 @@ func (t *TunnelAllWorkspaceUpdatesController) New(client WorkspaceUpdatesClient)
922930
updateHandler: t.updateHandler,
923931
ownerUsername: t.ownerUsername,
924932
recvLoopDone: make(chan struct{}),
925-
workspaces: make(map[uuid.UUID]*workspace),
933+
workspaces: make(map[uuid.UUID]*Workspace),
926934
}
927935
t.updater = updater
928936
go t.updater.recvLoop()
929937
return t.updater
930938
}
931939

932-
func (t *TunnelAllWorkspaceUpdatesController) CurrentState() *proto.WorkspaceUpdate {
940+
func (t *TunnelAllWorkspaceUpdatesController) CurrentState() (WorkspaceUpdate, error) {
933941
t.mu.Lock()
934942
defer t.mu.Unlock()
935943
if t.updater == nil {
936-
return nil
944+
return WorkspaceUpdate{}, xerrors.New("no updater")
937945
}
938946
t.updater.Lock()
939947
defer t.updater.Unlock()
940-
out := &proto.WorkspaceUpdate{
941-
UpsertedWorkspaces: make([]*proto.Workspace, 0, len(t.updater.workspaces)),
942-
UpsertedAgents: make([]*proto.Agent, 0, len(t.updater.workspaces)),
948+
out := WorkspaceUpdate{
949+
UpsertedWorkspaces: make([]*Workspace, 0, len(t.updater.workspaces)),
950+
UpsertedAgents: make([]*Agent, 0, len(t.updater.workspaces)),
951+
DeletedWorkspaces: make([]*Workspace, 0),
952+
DeletedAgents: make([]*Agent, 0),
943953
}
944954
for _, w := range t.updater.workspaces {
945-
upw := &proto.Workspace{
946-
Id: UUIDToByteSlice(w.id),
947-
Name: w.name,
948-
}
949-
out.UpsertedWorkspaces = append(out.UpsertedWorkspaces, upw)
955+
out.UpsertedWorkspaces = append(out.UpsertedWorkspaces, &Workspace{
956+
ID: w.ID,
957+
Name: w.Name,
958+
Status: w.Status,
959+
})
950960
for _, a := range w.agents {
951-
out.UpsertedAgents = append(out.UpsertedAgents, &proto.Agent{
952-
Id: UUIDToByteSlice(a.id),
953-
Name: a.name,
954-
WorkspaceId: UUIDToByteSlice(w.id),
961+
out.UpsertedAgents = append(out.UpsertedAgents, &Agent{
962+
ID: a.ID,
963+
Name: a.Name,
964+
WorkspaceID: a.WorkspaceID,
965+
Hosts: maps.Clone(a.Hosts),
955966
})
956967
}
957968
}
958-
return out
969+
return out, nil
959970
}
960971

961972
type tunnelUpdater struct {
@@ -969,7 +980,7 @@ type tunnelUpdater struct {
969980
recvLoopDone chan struct{}
970981

971982
sync.Mutex
972-
workspaces map[uuid.UUID]*workspace
983+
workspaces map[uuid.UUID]*Workspace
973984
closed bool
974985
}
975986

@@ -1031,20 +1042,78 @@ func (t *tunnelUpdater) recvLoop() {
10311042
}
10321043
}
10331044

1045+
type WorkspaceUpdate struct {
1046+
UpsertedWorkspaces []*Workspace
1047+
UpsertedAgents []*Agent
1048+
DeletedWorkspaces []*Workspace
1049+
DeletedAgents []*Agent
1050+
}
1051+
1052+
func (w *WorkspaceUpdate) Clone() WorkspaceUpdate {
1053+
clone := WorkspaceUpdate{
1054+
UpsertedWorkspaces: make([]*Workspace, len(w.UpsertedWorkspaces)),
1055+
UpsertedAgents: make([]*Agent, len(w.UpsertedAgents)),
1056+
DeletedWorkspaces: make([]*Workspace, len(w.DeletedWorkspaces)),
1057+
DeletedAgents: make([]*Agent, len(w.DeletedAgents)),
1058+
}
1059+
for i, ws := range w.UpsertedWorkspaces {
1060+
clone.UpsertedWorkspaces[i] = &Workspace{
1061+
ID: ws.ID,
1062+
Name: ws.Name,
1063+
Status: ws.Status,
1064+
}
1065+
}
1066+
for i, a := range w.UpsertedAgents {
1067+
clone.UpsertedAgents[i] = &Agent{
1068+
ID: a.ID,
1069+
Name: a.Name,
1070+
WorkspaceID: a.WorkspaceID,
1071+
Hosts: maps.Clone(a.Hosts),
1072+
}
1073+
}
1074+
for i, ws := range w.DeletedWorkspaces {
1075+
clone.DeletedWorkspaces[i] = &Workspace{
1076+
ID: ws.ID,
1077+
Name: ws.Name,
1078+
Status: ws.Status,
1079+
}
1080+
}
1081+
for i, a := range w.DeletedAgents {
1082+
clone.DeletedAgents[i] = &Agent{
1083+
ID: a.ID,
1084+
Name: a.Name,
1085+
WorkspaceID: a.WorkspaceID,
1086+
Hosts: maps.Clone(a.Hosts),
1087+
}
1088+
}
1089+
return clone
1090+
}
1091+
10341092
func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
10351093
t.Lock()
10361094
defer t.Unlock()
1095+
1096+
currentUpdate := WorkspaceUpdate{
1097+
UpsertedWorkspaces: []*Workspace{},
1098+
UpsertedAgents: []*Agent{},
1099+
DeletedWorkspaces: []*Workspace{},
1100+
DeletedAgents: []*Agent{},
1101+
}
1102+
10371103
for _, uw := range update.UpsertedWorkspaces {
10381104
workspaceID, err := uuid.FromBytes(uw.Id)
10391105
if err != nil {
10401106
return xerrors.Errorf("failed to parse workspace ID: %w", err)
10411107
}
1042-
w := workspace{
1043-
id: workspaceID,
1044-
name: uw.Name,
1045-
agents: make(map[uuid.UUID]agent),
1108+
w := &Workspace{
1109+
ID: workspaceID,
1110+
Name: uw.Name,
1111+
Status: uw.Status,
1112+
ownerUsername: t.ownerUsername,
1113+
agents: make(map[uuid.UUID]*Agent),
10461114
}
10471115
t.upsertWorkspace(w)
1116+
currentUpdate.UpsertedWorkspaces = append(currentUpdate.UpsertedWorkspaces, w)
10481117
}
10491118

10501119
// delete agents before deleting workspaces, since the agents have workspace ID references
@@ -1057,17 +1126,22 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
10571126
if err != nil {
10581127
return xerrors.Errorf("failed to parse workspace ID: %w", err)
10591128
}
1060-
err = t.deleteAgent(workspaceID, agentID)
1129+
deletedAgent, err := t.deleteAgent(workspaceID, agentID)
10611130
if err != nil {
10621131
return xerrors.Errorf("failed to delete agent: %w", err)
10631132
}
1133+
currentUpdate.DeletedAgents = append(currentUpdate.DeletedAgents, deletedAgent)
10641134
}
10651135
for _, dw := range update.DeletedWorkspaces {
10661136
workspaceID, err := uuid.FromBytes(dw.Id)
10671137
if err != nil {
10681138
return xerrors.Errorf("failed to parse workspace ID: %w", err)
10691139
}
1070-
t.deleteWorkspace(workspaceID)
1140+
deletedWorkspace, err := t.deleteWorkspace(workspaceID)
1141+
if err != nil {
1142+
return xerrors.Errorf("failed to delete workspace: %w", err)
1143+
}
1144+
currentUpdate.DeletedWorkspaces = append(currentUpdate.DeletedWorkspaces, deletedWorkspace)
10711145
}
10721146

10731147
// upsert agents last, after all workspaces have been added and deleted, since agents reference
@@ -1081,17 +1155,18 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
10811155
if err != nil {
10821156
return xerrors.Errorf("failed to parse workspace ID: %w", err)
10831157
}
1084-
a := agent{name: ua.Name, id: agentID}
1158+
a := &Agent{Name: ua.Name, ID: agentID, WorkspaceID: workspaceID}
10851159
err = t.upsertAgent(workspaceID, a)
10861160
if err != nil {
10871161
return xerrors.Errorf("failed to upsert agent: %w", err)
10881162
}
1163+
currentUpdate.UpsertedAgents = append(currentUpdate.UpsertedAgents, a)
10891164
}
10901165
allAgents := t.allAgentIDs()
10911166
t.coordCtrl.SyncDestinations(allAgents)
1167+
dnsNames := t.updateDNSNames()
10921168
if t.dnsHostsSetter != nil {
10931169
t.logger.Debug(context.Background(), "updating dns hosts")
1094-
dnsNames := t.allDNSNames()
10951170
err := t.dnsHostsSetter.SetDNSHosts(dnsNames)
10961171
if err != nil {
10971172
return xerrors.Errorf("failed to set DNS hosts: %w", err)
@@ -1101,43 +1176,55 @@ func (t *tunnelUpdater) handleUpdate(update *proto.WorkspaceUpdate) error {
11011176
}
11021177
if t.updateHandler != nil {
11031178
t.logger.Debug(context.Background(), "calling update handler")
1104-
err := t.updateHandler.Update(update)
1179+
err := t.updateHandler.Update(currentUpdate.Clone())
11051180
if err != nil {
11061181
t.logger.Error(context.Background(), "failed to call update handler", slog.Error(err))
11071182
}
11081183
}
11091184
return nil
11101185
}
11111186

1112-
func (t *tunnelUpdater) upsertWorkspace(w workspace) {
1113-
old, ok := t.workspaces[w.id]
1187+
func (t *tunnelUpdater) upsertWorkspace(w *Workspace) *Workspace {
1188+
old, ok := t.workspaces[w.ID]
11141189
if !ok {
1115-
t.workspaces[w.id] = &w
1116-
return
1190+
t.workspaces[w.ID] = w
1191+
return w
11171192
}
1118-
old.name = w.name
1193+
old.Name = w.Name
1194+
old.Status = w.Status
1195+
old.ownerUsername = w.ownerUsername
1196+
return w
11191197
}
11201198

1121-
func (t *tunnelUpdater) deleteWorkspace(id uuid.UUID) {
1199+
func (t *tunnelUpdater) deleteWorkspace(id uuid.UUID) (*Workspace, error) {
1200+
w, ok := t.workspaces[id]
1201+
if !ok {
1202+
return nil, xerrors.Errorf("workspace %s not found", id)
1203+
}
11221204
delete(t.workspaces, id)
1205+
return w, nil
11231206
}
11241207

1125-
func (t *tunnelUpdater) upsertAgent(workspaceID uuid.UUID, a agent) error {
1208+
func (t *tunnelUpdater) upsertAgent(workspaceID uuid.UUID, a *Agent) error {
11261209
w, ok := t.workspaces[workspaceID]
11271210
if !ok {
11281211
return xerrors.Errorf("workspace %s not found", workspaceID)
11291212
}
1130-
w.agents[a.id] = a
1213+
w.agents[a.ID] = a
11311214
return nil
11321215
}
11331216

1134-
func (t *tunnelUpdater) deleteAgent(workspaceID, id uuid.UUID) error {
1217+
func (t *tunnelUpdater) deleteAgent(workspaceID, id uuid.UUID) (*Agent, error) {
11351218
w, ok := t.workspaces[workspaceID]
11361219
if !ok {
1137-
return xerrors.Errorf("workspace %s not found", workspaceID)
1220+
return nil, xerrors.Errorf("workspace %s not found", workspaceID)
1221+
}
1222+
a, ok := w.agents[id]
1223+
if !ok {
1224+
return nil, xerrors.Errorf("agent %s not found in workspace %s", id, workspaceID)
11381225
}
11391226
delete(w.agents, id)
1140-
return nil
1227+
return a, nil
11411228
}
11421229

11431230
func (t *tunnelUpdater) allAgentIDs() []uuid.UUID {
@@ -1150,19 +1237,25 @@ func (t *tunnelUpdater) allAgentIDs() []uuid.UUID {
11501237
return out
11511238
}
11521239

1153-
func (t *tunnelUpdater) allDNSNames() map[dnsname.FQDN][]netip.Addr {
1240+
// updateDNSNames updates the DNS names for all workspaces in the tunnelUpdater.
1241+
func (t *tunnelUpdater) updateDNSNames() map[dnsname.FQDN][]netip.Addr {
11541242
names := make(map[dnsname.FQDN][]netip.Addr)
11551243
for _, w := range t.workspaces {
1156-
err := w.addAllDNSNames(names, t.ownerUsername)
1244+
err := w.updateDNSNames()
11571245
if err != nil {
11581246
// This should never happen in production, because converting the FQDN only fails
11591247
// if names are too long, and we put strict length limits on agent, workspace, and user
11601248
// names.
11611249
t.logger.Critical(context.Background(),
11621250
"failed to include DNS name(s)",
1163-
slog.F("workspace_id", w.id),
1251+
slog.F("workspace_id", w.ID),
11641252
slog.Error(err))
11651253
}
1254+
for _, a := range w.agents {
1255+
for name, addrs := range a.Hosts {
1256+
names[name] = addrs
1257+
}
1258+
}
11661259
}
11671260
return names
11681261
}

0 commit comments

Comments
 (0)