feat: improve coordinator peer mapping performance (#23696)

- Skipping DB querying entirely for peers that aren't actually connected
to our coordinator
- Opportunistically batching the queries for peers
This commit is contained in:
Jon Ayers
2026-04-03 14:22:58 -05:00
committed by GitHub
parent 01b8cdb00d
commit 333503f74e
10 changed files with 352 additions and 199 deletions
+4 -4
View File
@@ -3646,18 +3646,18 @@ func (q *querier) GetTailnetPeers(ctx context.Context, id uuid.UUID) ([]database
return q.db.GetTailnetPeers(ctx, id)
}
func (q *querier) GetTailnetTunnelPeerBindings(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsRow, error) {
func (q *querier) GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsBatchRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTailnetCoordinator); err != nil {
return nil, err
}
return q.db.GetTailnetTunnelPeerBindings(ctx, srcID)
return q.db.GetTailnetTunnelPeerBindingsBatch(ctx, ids)
}
func (q *querier) GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerIDsRow, error) {
func (q *querier) GetTailnetTunnelPeerIDsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerIDsBatchRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTailnetCoordinator); err != nil {
return nil, err
}
return q.db.GetTailnetTunnelPeerIDs(ctx, srcID)
return q.db.GetTailnetTunnelPeerIDsBatch(ctx, ids)
}
func (q *querier) GetTaskByID(ctx context.Context, id uuid.UUID) (database.Task, error) {
+4 -6
View File
@@ -3750,13 +3750,11 @@ func (s *MethodTestSuite) TestTailnetFunctions() {
check.Args(uuid.New()).
Asserts(rbac.ResourceTailnetCoordinator, policy.ActionRead)
}))
s.Run("GetTailnetTunnelPeerBindings", s.Subtest(func(_ database.Store, check *expects) {
check.Args(uuid.New()).
Asserts(rbac.ResourceTailnetCoordinator, policy.ActionRead)
s.Run("GetTailnetTunnelPeerBindingsBatch", s.Subtest(func(_ database.Store, check *expects) {
check.Args([]uuid.UUID{uuid.New()}).Asserts(rbac.ResourceTailnetCoordinator, policy.ActionRead)
}))
s.Run("GetTailnetTunnelPeerIDs", s.Subtest(func(_ database.Store, check *expects) {
check.Args(uuid.New()).
Asserts(rbac.ResourceTailnetCoordinator, policy.ActionRead)
s.Run("GetTailnetTunnelPeerIDsBatch", s.Subtest(func(_ database.Store, check *expects) {
check.Args([]uuid.UUID{uuid.New()}).Asserts(rbac.ResourceTailnetCoordinator, policy.ActionRead)
}))
s.Run("GetAllTailnetCoordinators", s.Subtest(func(_ database.Store, check *expects) {
check.Args().
+8 -8
View File
@@ -2216,19 +2216,19 @@ func (m queryMetricsStore) GetTailnetPeers(ctx context.Context, id uuid.UUID) ([
return r0, r1
}
func (m queryMetricsStore) GetTailnetTunnelPeerBindings(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsRow, error) {
func (m queryMetricsStore) GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsBatchRow, error) {
start := time.Now()
r0, r1 := m.s.GetTailnetTunnelPeerBindings(ctx, srcID)
m.queryLatencies.WithLabelValues("GetTailnetTunnelPeerBindings").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTailnetTunnelPeerBindings").Inc()
r0, r1 := m.s.GetTailnetTunnelPeerBindingsBatch(ctx, ids)
m.queryLatencies.WithLabelValues("GetTailnetTunnelPeerBindingsBatch").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTailnetTunnelPeerBindingsBatch").Inc()
return r0, r1
}
func (m queryMetricsStore) GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerIDsRow, error) {
func (m queryMetricsStore) GetTailnetTunnelPeerIDsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerIDsBatchRow, error) {
start := time.Now()
r0, r1 := m.s.GetTailnetTunnelPeerIDs(ctx, srcID)
m.queryLatencies.WithLabelValues("GetTailnetTunnelPeerIDs").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTailnetTunnelPeerIDs").Inc()
r0, r1 := m.s.GetTailnetTunnelPeerIDsBatch(ctx, ids)
m.queryLatencies.WithLabelValues("GetTailnetTunnelPeerIDsBatch").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTailnetTunnelPeerIDsBatch").Inc()
return r0, r1
}
+14 -14
View File
@@ -4113,34 +4113,34 @@ func (mr *MockStoreMockRecorder) GetTailnetPeers(ctx, id any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTailnetPeers", reflect.TypeOf((*MockStore)(nil).GetTailnetPeers), ctx, id)
}
// GetTailnetTunnelPeerBindings mocks base method.
func (m *MockStore) GetTailnetTunnelPeerBindings(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsRow, error) {
// GetTailnetTunnelPeerBindingsBatch mocks base method.
func (m *MockStore) GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerBindingsBatchRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetTailnetTunnelPeerBindings", ctx, srcID)
ret0, _ := ret[0].([]database.GetTailnetTunnelPeerBindingsRow)
ret := m.ctrl.Call(m, "GetTailnetTunnelPeerBindingsBatch", ctx, ids)
ret0, _ := ret[0].([]database.GetTailnetTunnelPeerBindingsBatchRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetTailnetTunnelPeerBindings indicates an expected call of GetTailnetTunnelPeerBindings.
func (mr *MockStoreMockRecorder) GetTailnetTunnelPeerBindings(ctx, srcID any) *gomock.Call {
// GetTailnetTunnelPeerBindingsBatch indicates an expected call of GetTailnetTunnelPeerBindingsBatch.
func (mr *MockStoreMockRecorder) GetTailnetTunnelPeerBindingsBatch(ctx, ids any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTailnetTunnelPeerBindings", reflect.TypeOf((*MockStore)(nil).GetTailnetTunnelPeerBindings), ctx, srcID)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTailnetTunnelPeerBindingsBatch", reflect.TypeOf((*MockStore)(nil).GetTailnetTunnelPeerBindingsBatch), ctx, ids)
}
// GetTailnetTunnelPeerIDs mocks base method.
func (m *MockStore) GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUID) ([]database.GetTailnetTunnelPeerIDsRow, error) {
// GetTailnetTunnelPeerIDsBatch mocks base method.
func (m *MockStore) GetTailnetTunnelPeerIDsBatch(ctx context.Context, ids []uuid.UUID) ([]database.GetTailnetTunnelPeerIDsBatchRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetTailnetTunnelPeerIDs", ctx, srcID)
ret0, _ := ret[0].([]database.GetTailnetTunnelPeerIDsRow)
ret := m.ctrl.Call(m, "GetTailnetTunnelPeerIDsBatch", ctx, ids)
ret0, _ := ret[0].([]database.GetTailnetTunnelPeerIDsBatchRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetTailnetTunnelPeerIDs indicates an expected call of GetTailnetTunnelPeerIDs.
func (mr *MockStoreMockRecorder) GetTailnetTunnelPeerIDs(ctx, srcID any) *gomock.Call {
// GetTailnetTunnelPeerIDsBatch indicates an expected call of GetTailnetTunnelPeerIDsBatch.
func (mr *MockStoreMockRecorder) GetTailnetTunnelPeerIDsBatch(ctx, ids any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTailnetTunnelPeerIDs", reflect.TypeOf((*MockStore)(nil).GetTailnetTunnelPeerIDs), ctx, srcID)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTailnetTunnelPeerIDsBatch", reflect.TypeOf((*MockStore)(nil).GetTailnetTunnelPeerIDsBatch), ctx, ids)
}
// GetTaskByID mocks base method.
+2 -2
View File
@@ -477,8 +477,8 @@ type sqlcQuerier interface {
// Used for recovery after coderd crashes or long hangs.
GetStaleChats(ctx context.Context, staleThreshold time.Time) ([]Chat, error)
GetTailnetPeers(ctx context.Context, id uuid.UUID) ([]TailnetPeer, error)
GetTailnetTunnelPeerBindings(ctx context.Context, srcID uuid.UUID) ([]GetTailnetTunnelPeerBindingsRow, error)
GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUID) ([]GetTailnetTunnelPeerIDsRow, error)
GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]GetTailnetTunnelPeerBindingsBatchRow, error)
GetTailnetTunnelPeerIDsBatch(ctx context.Context, ids []uuid.UUID) ([]GetTailnetTunnelPeerIDsBatchRow, error)
GetTaskByID(ctx context.Context, id uuid.UUID) (Task, error)
GetTaskByOwnerIDAndName(ctx context.Context, arg GetTaskByOwnerIDAndNameParams) (Task, error)
GetTaskByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) (Task, error)
+35 -30
View File
@@ -19309,43 +19309,44 @@ func (q *sqlQuerier) GetTailnetPeers(ctx context.Context, id uuid.UUID) ([]Tailn
return items, nil
}
const getTailnetTunnelPeerBindings = `-- name: GetTailnetTunnelPeerBindings :many
SELECT id AS peer_id, coordinator_id, updated_at, node, status
FROM tailnet_peers
WHERE id IN (
SELECT dst_id as peer_id
FROM tailnet_tunnels
WHERE tailnet_tunnels.src_id = $1
const getTailnetTunnelPeerBindingsBatch = `-- name: GetTailnetTunnelPeerBindingsBatch :many
SELECT tp.id AS peer_id, tp.coordinator_id, tp.updated_at, tp.node, tp.status,
tunnels.lookup_id
FROM (
SELECT dst_id AS peer_id, src_id AS lookup_id
FROM tailnet_tunnels WHERE src_id = ANY($1 :: uuid[])
UNION
SELECT src_id as peer_id
FROM tailnet_tunnels
WHERE tailnet_tunnels.dst_id = $1
)
SELECT src_id AS peer_id, dst_id AS lookup_id
FROM tailnet_tunnels WHERE dst_id = ANY($1 :: uuid[])
) tunnels
INNER JOIN tailnet_peers tp ON tp.id = tunnels.peer_id
`
type GetTailnetTunnelPeerBindingsRow struct {
type GetTailnetTunnelPeerBindingsBatchRow struct {
PeerID uuid.UUID `db:"peer_id" json:"peer_id"`
CoordinatorID uuid.UUID `db:"coordinator_id" json:"coordinator_id"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
Node []byte `db:"node" json:"node"`
Status TailnetStatus `db:"status" json:"status"`
LookupID uuid.UUID `db:"lookup_id" json:"lookup_id"`
}
func (q *sqlQuerier) GetTailnetTunnelPeerBindings(ctx context.Context, srcID uuid.UUID) ([]GetTailnetTunnelPeerBindingsRow, error) {
rows, err := q.db.QueryContext(ctx, getTailnetTunnelPeerBindings, srcID)
func (q *sqlQuerier) GetTailnetTunnelPeerBindingsBatch(ctx context.Context, ids []uuid.UUID) ([]GetTailnetTunnelPeerBindingsBatchRow, error) {
rows, err := q.db.QueryContext(ctx, getTailnetTunnelPeerBindingsBatch, pq.Array(ids))
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetTailnetTunnelPeerBindingsRow
var items []GetTailnetTunnelPeerBindingsBatchRow
for rows.Next() {
var i GetTailnetTunnelPeerBindingsRow
var i GetTailnetTunnelPeerBindingsBatchRow
if err := rows.Scan(
&i.PeerID,
&i.CoordinatorID,
&i.UpdatedAt,
&i.Node,
&i.Status,
&i.LookupID,
); err != nil {
return nil, err
}
@@ -19360,32 +19361,36 @@ func (q *sqlQuerier) GetTailnetTunnelPeerBindings(ctx context.Context, srcID uui
return items, nil
}
const getTailnetTunnelPeerIDs = `-- name: GetTailnetTunnelPeerIDs :many
SELECT dst_id as peer_id, coordinator_id, updated_at
FROM tailnet_tunnels
WHERE tailnet_tunnels.src_id = $1
UNION
SELECT src_id as peer_id, coordinator_id, updated_at
FROM tailnet_tunnels
WHERE tailnet_tunnels.dst_id = $1
const getTailnetTunnelPeerIDsBatch = `-- name: GetTailnetTunnelPeerIDsBatch :many
SELECT src_id AS lookup_id, dst_id AS peer_id, coordinator_id, updated_at
FROM tailnet_tunnels WHERE src_id = ANY($1 :: uuid[])
UNION ALL
SELECT dst_id AS lookup_id, src_id AS peer_id, coordinator_id, updated_at
FROM tailnet_tunnels WHERE dst_id = ANY($1 :: uuid[])
`
type GetTailnetTunnelPeerIDsRow struct {
type GetTailnetTunnelPeerIDsBatchRow struct {
LookupID uuid.UUID `db:"lookup_id" json:"lookup_id"`
PeerID uuid.UUID `db:"peer_id" json:"peer_id"`
CoordinatorID uuid.UUID `db:"coordinator_id" json:"coordinator_id"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
}
func (q *sqlQuerier) GetTailnetTunnelPeerIDs(ctx context.Context, srcID uuid.UUID) ([]GetTailnetTunnelPeerIDsRow, error) {
rows, err := q.db.QueryContext(ctx, getTailnetTunnelPeerIDs, srcID)
func (q *sqlQuerier) GetTailnetTunnelPeerIDsBatch(ctx context.Context, ids []uuid.UUID) ([]GetTailnetTunnelPeerIDsBatchRow, error) {
rows, err := q.db.QueryContext(ctx, getTailnetTunnelPeerIDsBatch, pq.Array(ids))
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetTailnetTunnelPeerIDsRow
var items []GetTailnetTunnelPeerIDsBatchRow
for rows.Next() {
var i GetTailnetTunnelPeerIDsRow
if err := rows.Scan(&i.PeerID, &i.CoordinatorID, &i.UpdatedAt); err != nil {
var i GetTailnetTunnelPeerIDsBatchRow
if err := rows.Scan(
&i.LookupID,
&i.PeerID,
&i.CoordinatorID,
&i.UpdatedAt,
); err != nil {
return nil, err
}
items = append(items, i)
+19 -22
View File
@@ -96,28 +96,6 @@ DELETE
FROM tailnet_tunnels
WHERE coordinator_id = $1 and src_id = $2;
-- name: GetTailnetTunnelPeerIDs :many
SELECT dst_id as peer_id, coordinator_id, updated_at
FROM tailnet_tunnels
WHERE tailnet_tunnels.src_id = $1
UNION
SELECT src_id as peer_id, coordinator_id, updated_at
FROM tailnet_tunnels
WHERE tailnet_tunnels.dst_id = $1;
-- name: GetTailnetTunnelPeerBindings :many
SELECT id AS peer_id, coordinator_id, updated_at, node, status
FROM tailnet_peers
WHERE id IN (
SELECT dst_id as peer_id
FROM tailnet_tunnels
WHERE tailnet_tunnels.src_id = $1
UNION
SELECT src_id as peer_id
FROM tailnet_tunnels
WHERE tailnet_tunnels.dst_id = $1
);
-- For PG Coordinator HTMLDebug
-- name: GetAllTailnetCoordinators :many
@@ -128,3 +106,22 @@ SELECT * FROM tailnet_peers;
-- name: GetAllTailnetTunnels :many
SELECT * FROM tailnet_tunnels;
-- name: GetTailnetTunnelPeerIDsBatch :many
SELECT src_id AS lookup_id, dst_id AS peer_id, coordinator_id, updated_at
FROM tailnet_tunnels WHERE src_id = ANY(@ids :: uuid[])
UNION ALL
SELECT dst_id AS lookup_id, src_id AS peer_id, coordinator_id, updated_at
FROM tailnet_tunnels WHERE dst_id = ANY(@ids :: uuid[]);
-- name: GetTailnetTunnelPeerBindingsBatch :many
SELECT tp.id AS peer_id, tp.coordinator_id, tp.updated_at, tp.node, tp.status,
tunnels.lookup_id
FROM (
SELECT dst_id AS peer_id, src_id AS lookup_id
FROM tailnet_tunnels WHERE src_id = ANY(@ids :: uuid[])
UNION
SELECT src_id AS peer_id, dst_id AS lookup_id
FROM tailnet_tunnels WHERE dst_id = ANY(@ids :: uuid[])
) tunnels
INNER JOIN tailnet_peers tp ON tp.id = tunnels.peer_id;