Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a1d51f0

Browse files
sreyasreya
authored
feat: batch connection logs to avoid DB lock contention (#23727)
- Running 30k connections was generating a ton of lock contention in the DB
1 parent 333503f commit a1d51f0

21 files changed

Lines changed: 2159 additions & 417 deletions

‎coderd/agentapi/connectionlog.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
8585
AgentName: a.AgentName,
8686
Type: connectionType,
8787
Code: code,
88-
Ip: logIP,
88+
IP: logIP,
8989
ConnectionID: uuid.NullUUID{
9090
UUID: connectionID,
9191
Valid: true,

‎coderd/agentapi/connectionlog_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ func TestConnectionLog(t *testing.T) {
152152
Int32: tt.status,
153153
Valid: *tt.action == agentproto.Connection_DISCONNECT,
154154
},
155-
Ip: expectedIP,
155+
IP: expectedIP,
156156
Type: agentProtoConnectionTypeToConnectionLog(t, *tt.typ),
157157
DisconnectReason: sql.NullString{
158158
String: tt.reason,

‎coderd/connectionlog/connectionlog.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@ func (m *FakeConnectionLogger) Contains(t testing.TB, expected database.UpsertCo
9090
t.Logf("connection log %d: expected Code %d, got %d", idx+1, expected.Code.Int32, cl.Code.Int32)
9191
continue
9292
}
93-
if expected.Ip.Valid && cl.Ip.IPNet.String() != expected.Ip.IPNet.String() {
94-
t.Logf("connection log %d: expected IP %s, got %s", idx+1, expected.Ip.IPNet, cl.Ip.IPNet)
93+
if expected.IP.Valid && cl.IP.IPNet.String() != expected.IP.IPNet.String() {
94+
t.Logf("connection log %d: expected IP %s, got %s", idx+1, expected.IP.IPNet, cl.IP.IPNet)
9595
continue
9696
}
9797
if expected.UserAgent.Valid && cl.UserAgent.String != expected.UserAgent.String {

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1627,6 +1627,13 @@ func (q *querier) BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg datab
16271627
return q.db.BatchUpdateWorkspaceNextStartAt(ctx, arg)
16281628
}
16291629

1630+
func (q *querier) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error {
1631+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
1632+
return err
1633+
}
1634+
return q.db.BatchUpsertConnectionLogs(ctx, arg)
1635+
}
1636+
16301637
func (q *querier) BulkMarkNotificationMessagesFailed(ctx context.Context, arg database.BulkMarkNotificationMessagesFailedParams) (int64, error) {
16311638
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceNotificationMessage); err != nil {
16321639
return 0, err
@@ -7065,13 +7072,6 @@ func (q *querier) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl strin
70657072
return q.db.UpsertChatWorkspaceTTL(ctx, workspaceTtl)
70667073
}
70677074

7068-
func (q *querier) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
7069-
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
7070-
return database.ConnectionLog{}, err
7071-
}
7072-
return q.db.UpsertConnectionLog(ctx, arg)
7073-
}
7074-
70757075
func (q *querier) UpsertDefaultProxy(ctx context.Context, arg database.UpsertDefaultProxyParams) error {
70767076
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
70777077
return err

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -338,10 +338,9 @@ func (s *MethodTestSuite) TestAuditLogs() {
338338
}
339339

340340
func (s *MethodTestSuite) TestConnectionLogs() {
341-
s.Run("UpsertConnectionLog", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
342-
ws := testutil.Fake(s.T(), faker, database.WorkspaceTable{})
343-
arg := database.UpsertConnectionLogParams{Ip: defaultIPAddress(), Type: database.ConnectionTypeSsh, WorkspaceID: ws.ID, OrganizationID: ws.OrganizationID, ConnectionStatus: database.ConnectionStatusConnected, WorkspaceOwnerID: ws.OwnerID}
344-
dbm.EXPECT().UpsertConnectionLog(gomock.Any(), arg).Return(database.ConnectionLog{}, nil).AnyTimes()
341+
s.Run("BatchUpsertConnectionLogs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
342+
arg := database.BatchUpsertConnectionLogsParams{}
343+
dbm.EXPECT().BatchUpsertConnectionLogs(gomock.Any(), arg).Return(nil).AnyTimes()
345344
check.Args(arg).Asserts(rbac.ResourceConnectionLog, policy.ActionUpdate)
346345
}))
347346
s.Run("GetConnectionLogsOffset", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {

‎coderd/database/dbgen/dbgen.go‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ func AuditLog(t testing.TB, db database.Store, seed database.AuditLog) database.
7676
}
7777

7878
func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnectionLogParams) database.ConnectionLog {
79-
log, err := db.UpsertConnectionLog(genCtx, database.UpsertConnectionLogParams{
79+
arg := database.UpsertConnectionLogParams{
8080
ID: takeFirst(seed.ID, uuid.New()),
8181
Time: takeFirst(seed.Time, dbtime.Now()),
8282
OrganizationID: takeFirst(seed.OrganizationID, uuid.New()),
@@ -89,7 +89,7 @@ func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnecti
8989
Int32: takeFirst(seed.Code.Int32, 0),
9090
Valid: takeFirst(seed.Code.Valid, false),
9191
},
92-
Ip: pqtype.Inet{
92+
IP: pqtype.Inet{
9393
IPNet: net.IPNet{
9494
IP: net.IPv4(127, 0, 0, 1),
9595
Mask: net.IPv4Mask(255, 255, 255, 255),
@@ -117,9 +117,53 @@ func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnecti
117117
Valid: takeFirst(seed.DisconnectReason.Valid, false),
118118
},
119119
ConnectionStatus: takeFirst(seed.ConnectionStatus, database.ConnectionStatusConnected),
120+
}
121+
122+
var disconnectTime sql.NullTime
123+
if arg.ConnectionStatus == database.ConnectionStatusDisconnected {
124+
disconnectTime = sql.NullTime{Time: arg.Time, Valid: true}
125+
}
126+
127+
err := db.BatchUpsertConnectionLogs(genCtx, database.BatchUpsertConnectionLogsParams{
128+
ID: []uuid.UUID{arg.ID},
129+
ConnectTime: []time.Time{arg.Time},
130+
OrganizationID: []uuid.UUID{arg.OrganizationID},
131+
WorkspaceOwnerID: []uuid.UUID{arg.WorkspaceOwnerID},
132+
WorkspaceID: []uuid.UUID{arg.WorkspaceID},
133+
WorkspaceName: []string{arg.WorkspaceName},
134+
AgentName: []string{arg.AgentName},
135+
Type: []database.ConnectionType{arg.Type},
136+
Code: []int32{arg.Code.Int32},
137+
CodeValid: []bool{arg.Code.Valid},
138+
Ip: []pqtype.Inet{arg.IP},
139+
UserAgent: []string{arg.UserAgent.String},
140+
UserID: []uuid.UUID{arg.UserID.UUID},
141+
SlugOrPort: []string{arg.SlugOrPort.String},
142+
ConnectionID: []uuid.UUID{arg.ConnectionID.UUID},
143+
DisconnectReason: []string{arg.DisconnectReason.String},
144+
DisconnectTime: []time.Time{disconnectTime.Time},
120145
})
121146
require.NoError(t, err, "insert connection log")
122-
return log
147+
148+
// Query back the actual row from the database. On upsert
149+
// conflict the DB keeps the original row's ID, so we can't
150+
// rely on arg.ID. Match on the conflict key for rows with a
151+
// connection_id, or by primary key for NULL connection_id.
152+
rows, err := db.GetConnectionLogsOffset(genCtx, database.GetConnectionLogsOffsetParams{})
153+
require.NoError(t, err, "query connection logs")
154+
for _, row := range rows {
155+
if arg.ConnectionID.Valid {
156+
if row.ConnectionLog.ConnectionID == arg.ConnectionID &&
157+
row.ConnectionLog.WorkspaceID == arg.WorkspaceID &&
158+
row.ConnectionLog.AgentName == arg.AgentName {
159+
return row.ConnectionLog
160+
}
161+
} else if row.ConnectionLog.ID == arg.ID {
162+
return row.ConnectionLog
163+
}
164+
}
165+
require.Failf(t, "connection log not found", "id=%s", arg.ID)
166+
return database.ConnectionLog{} // unreachable
123167
}
124168

125169
func Template(t testing.TB, db database.Store, seed database.Template) database.Template {

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 8 additions & 8 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 14 additions & 15 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/modelmethods.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"time"
1111

1212
"github.com/google/uuid"
13+
"github.com/sqlc-dev/pqtype"
1314
"golang.org/x/exp/maps"
1415
"golang.org/x/oauth2"
1516
"golang.org/x/xerrors"
@@ -923,3 +924,28 @@ func WorkspaceIdentityFromWorkspace(w Workspace) WorkspaceIdentity {
923924
func (r GetWorkspaceAgentAndWorkspaceByIDRow) RBACObject() rbac.Object {
924925
return r.WorkspaceTable.RBACObject()
925926
}
927+
928+
// UpsertConnectionLogParams contains the parameters for upserting a
929+
// connection log entry. This struct is hand-maintained (not generated
930+
// by sqlc) because the single-row UpsertConnectionLog query was
931+
// removed in favor of BatchUpsertConnectionLogs, but the struct is
932+
// still used as the canonical connection log event type throughout
933+
// the codebase.
934+
type UpsertConnectionLogParams struct {
935+
ID uuid.UUID `db:"id" json:"id"`
936+
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
937+
WorkspaceOwnerID uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
938+
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
939+
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
940+
AgentName string `db:"agent_name" json:"agent_name"`
941+
Type ConnectionType `db:"type" json:"type"`
942+
Code sql.NullInt32 `db:"code" json:"code"`
943+
IP pqtype.Inet `db:"ip" json:"ip"`
944+
UserAgent sql.NullString `db:"user_agent" json:"user_agent"`
945+
UserID uuid.NullUUID `db:"user_id" json:"user_id"`
946+
SlugOrPort sql.NullString `db:"slug_or_port" json:"slug_or_port"`
947+
ConnectionID uuid.NullUUID `db:"connection_id" json:"connection_id"`
948+
DisconnectReason sql.NullString `db:"disconnect_reason" json:"disconnect_reason"`
949+
Time time.Time `db:"time" json:"time"`
950+
ConnectionStatus ConnectionStatus `db:"connection_status" json:"connection_status"`
951+
}

‎coderd/database/querier.go‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)