Skip to content

Commit 28ae155

Browse files
deansheatherdeansheather
committed
tests and various fixes
1 parent 2d2f1a3 commit 28ae155

File tree

7 files changed

+311
-39
lines changed

7 files changed

+311
-39
lines changed

‎coderd/database/dbfake/databasefake.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5210,10 +5210,14 @@ func (q *fakeQuerier) InsertWorkspaceProxy(_ context.Context, arg database.Inser
52105210
q.mutex.Lock()
52115211
defer q.mutex.Unlock()
52125212

5213+
lastRegionID := int32(0)
52135214
for _, p := range q.workspaceProxies {
52145215
if !p.Deleted && p.Name == arg.Name {
52155216
return database.WorkspaceProxy{}, errDuplicateKey
52165217
}
5218+
if p.RegionID > lastRegionID {
5219+
lastRegionID = p.RegionID
5220+
}
52175221
}
52185222

52195223
p := database.WorkspaceProxy{
@@ -5223,6 +5227,7 @@ func (q *fakeQuerier) InsertWorkspaceProxy(_ context.Context, arg database.Inser
52235227
Icon: arg.Icon,
52245228
DerpEnabled: arg.DerpEnabled,
52255229
TokenHashedSecret: arg.TokenHashedSecret,
5230+
RegionID: lastRegionID + 1,
52265231
CreatedAt: arg.CreatedAt,
52275232
UpdatedAt: arg.UpdatedAt,
52285233
Deleted: false,

‎codersdk/workspaceagents.go

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -147,29 +147,41 @@ type WorkspaceAgentConnectionInfo struct {
147147
DERPMap *tailcfg.DERPMap `json:"derp_map"`
148148
}
149149

150+
func (c *Client) WorkspaceAgentConnectionInfo(ctx context.Context, agentID uuid.UUID) (WorkspaceAgentConnectionInfo, error) {
151+
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspaceagents/%s/connection", agentID), nil)
152+
if err != nil {
153+
return WorkspaceAgentConnectionInfo{}, err
154+
}
155+
defer res.Body.Close()
156+
if res.StatusCode != http.StatusOK {
157+
return WorkspaceAgentConnectionInfo{}, ReadBodyAsError(res)
158+
}
159+
160+
var connInfo WorkspaceAgentConnectionInfo
161+
return connInfo, json.NewDecoder(res.Body).Decode(&connInfo)
162+
}
163+
150164
// @typescript-ignore DialWorkspaceAgentOptions
151165
type DialWorkspaceAgentOptions struct {
152166
Logger slog.Logger
153167
// BlockEndpoints forced a direct connection through DERP.
154168
BlockEndpoints bool
169+
// CustomConnectionInfo avoids hitting the API to get connection info.
170+
CustomConnectionInfo *WorkspaceAgentConnectionInfo
155171
}
156172

157173
func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, options *DialWorkspaceAgentOptions) (agentConn *WorkspaceAgentConn, err error) {
158174
if options == nil {
159175
options = &DialWorkspaceAgentOptions{}
160176
}
161-
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspaceagents/%s/connection", agentID), nil)
162-
if err != nil {
163-
return nil, err
164-
}
165-
defer res.Body.Close()
166-
if res.StatusCode != http.StatusOK {
167-
return nil, ReadBodyAsError(res)
168-
}
169-
var connInfo WorkspaceAgentConnectionInfo
170-
err = json.NewDecoder(res.Body).Decode(&connInfo)
171-
if err != nil {
172-
return nil, xerrors.Errorf("decode conn info: %w", err)
177+
178+
connInfo := options.CustomConnectionInfo
179+
if connInfo == nil {
180+
res, err := c.WorkspaceAgentConnectionInfo(ctx, agentID)
181+
if err != nil {
182+
return nil, xerrors.Errorf("get connection info: %w", err)
183+
}
184+
connInfo = &res
173185
}
174186

175187
ip := tailnet.IP()

‎enterprise/coderd/coderd.go

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,35 @@ func (api *API) updateEntitlements(ctx context.Context) error {
450450
return nil
451451
}
452452

453+
// getProxyDERPStartingRegionID returns the starting region ID that should be
454+
// used for workspace proxies. A proxy's actual region ID is the return value
455+
// from this function + it's RegionID field.
456+
//
457+
// Two ints are returned, the first is the starting region ID for proxies, and
458+
// the second is the maximum region ID that already exists in the DERP map.
459+
func getProxyDERPStartingRegionID(derpMap *tailcfg.DERPMap) (sID int, mID int) {
460+
maxRegionID := 0
461+
for _, region := range derpMap.Regions {
462+
if region.RegionID > maxRegionID {
463+
maxRegionID = region.RegionID
464+
}
465+
}
466+
if maxRegionID < 0 {
467+
maxRegionID = 0
468+
}
469+
470+
// Round to the nearest 10,000 with a sufficient buffer of at least 2,000.
471+
const roundStartingRegionID = 10_000
472+
const startingRegionIDBuffer = 2_000
473+
startingRegionID := maxRegionID + startingRegionIDBuffer
474+
startingRegionID = int(math.Ceil(float64(startingRegionID)/roundStartingRegionID) * roundStartingRegionID)
475+
if startingRegionID < roundStartingRegionID {
476+
startingRegionID = roundStartingRegionID
477+
}
478+
479+
return startingRegionID, maxRegionID
480+
}
481+
453482
var (
454483
lastDerpConflictMutex sync.Mutex
455484
lastDerpConflictLog time.Time
@@ -461,16 +490,8 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
461490

462491
// Find the starting region ID that we'll use for proxies. This must be
463492
// deterministic based on the derp map.
464-
startingRegionID := 0
465-
for _, region := range derpMap.Regions {
466-
if region.RegionID > startingRegionID {
467-
startingRegionID = region.RegionID
468-
}
469-
}
470-
if startingRegionID < 0 {
471-
startingRegionID = 0
472-
}
473-
if startingRegionID >= 1<<32 {
493+
startingRegionID, largestRegionID := getProxyDERPStartingRegionID(derpMap)
494+
if largestRegionID >= 1<<32 {
474495
// Enforce an upper bound on the region ID. This shouldn't be hit in
475496
// practice, but it's a good sanity check.
476497
lastDerpConflictMutex.Lock()
@@ -483,23 +504,13 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
483504
logger.Warn(
484505
context.Background(),
485506
"existing DERP region IDs are too large, proxy region IDs will not be populated in the derp map. Please ensure that all DERP region IDs are less than 2^32.",
486-
slog.F("starting_region_id", startingRegionID),
507+
slog.F("largest_region_id", largestRegionID),
487508
slog.F("max_region_id", 1<<32-1),
488509
)
489510
return derpMap
490511
}
491512
}
492513

493-
// Round to the nearest 10,000 with a sufficient buffer of at least
494-
// 2,000.
495-
const roundStartingRegionID = 10_000
496-
const startingRegionIDBuffer = 2_000
497-
startingRegionID += startingRegionIDBuffer
498-
startingRegionID = int(math.Ceil(float64(startingRegionID)/roundStartingRegionID) * roundStartingRegionID)
499-
if startingRegionID < roundStartingRegionID {
500-
startingRegionID = roundStartingRegionID
501-
}
502-
503514
// Add all healthy proxies to the DERP map.
504515
statusMap := proxyHealth.HealthStatus()
505516
statusLoop:
@@ -564,12 +575,13 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
564575
}
565576

566577
derpMap.Regions[regionID] = &tailcfg.DERPRegion{
567-
EmbeddedRelay: true,
578+
// EmbeddedRelay ONLY applies to the primary.
579+
EmbeddedRelay: false,
568580
RegionID: regionID,
569581
RegionCode: regionCode,
570582
RegionName: status.Proxy.Name,
571583
Nodes: []*tailcfg.DERPNode{{
572-
Name: fmt.Sprintf("%db", regionID),
584+
Name: fmt.Sprintf("%da", regionID),
573585
RegionID: regionID,
574586
HostName: u.Hostname(),
575587
DERPPort: portInt,

‎enterprise/coderd/coderdenttest/proxytest.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,9 @@ func NewWorkspaceProxy(t *testing.T, coderdAPI *coderd.API, owner *codersdk.Clie
129129
DisablePathApps: options.DisablePathApps,
130130
// We need a new registry to not conflict with the coderd internal
131131
// proxy metrics.
132-
PrometheusRegistry: prometheus.NewRegistry(),
132+
PrometheusRegistry: prometheus.NewRegistry(),
133+
DERPEnabled: true,
134+
DERPServerRelayAddress: accessURL.String(),
133135
})
134136
require.NoError(t, err)
135137
t.Cleanup(func() {

‎enterprise/coderd/workspaceproxy.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -374,8 +374,8 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
374374
return
375375
}
376376

377-
// TODO: get region ID
378-
var regionID int32 = 1234
377+
startingRegionID, _ := getProxyDERPStartingRegionID(api.Options.DERPMap)
378+
regionID := int32(startingRegionID) + proxy.RegionID
379379

380380
err := api.Database.InTx(func(db database.Store) error {
381381
// First, update the proxy's values in the database.
@@ -473,6 +473,7 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
473473
httpapi.Write(ctx, rw, http.StatusCreated, wsproxysdk.RegisterWorkspaceProxyResponse{
474474
AppSecurityKey: api.AppSecurityKey.String(),
475475
DERPMeshKey: api.DERPServer.MeshKey(),
476+
DERPRegionID: regionID,
476477
SiblingReplicas: siblingsRes,
477478
})
478479

0 commit comments

Comments
 (0)