mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
feat: batch connection logs to avoid DB lock contention (#23727)
- Running 30k connections was generating a ton of lock contention in the DB
This commit is contained in:
Jon Ayers
+117
-114
@@ -7338,6 +7338,123 @@ func (q *sqlQuerier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg U
|
||||
return i, err
|
||||
}
|
||||
|
||||
const batchUpsertConnectionLogs = `-- name: BatchUpsertConnectionLogs :exec
|
||||
INSERT INTO connection_logs (
|
||||
id, connect_time, organization_id, workspace_owner_id, workspace_id,
|
||||
workspace_name, agent_name, type, code, ip, user_agent, user_id,
|
||||
slug_or_port, connection_id, disconnect_reason, disconnect_time
|
||||
)
|
||||
SELECT
|
||||
u.id,
|
||||
u.connect_time,
|
||||
u.organization_id,
|
||||
u.workspace_owner_id,
|
||||
u.workspace_id,
|
||||
u.workspace_name,
|
||||
u.agent_name,
|
||||
u.type,
|
||||
-- Use the validity flag to distinguish "no code" (NULL) from a
|
||||
-- legitimate zero exit code.
|
||||
CASE WHEN u.code_valid THEN u.code ELSE NULL END,
|
||||
u.ip,
|
||||
NULLIF(u.user_agent, ''),
|
||||
NULLIF(u.user_id, '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
NULLIF(u.slug_or_port, ''),
|
||||
NULLIF(u.connection_id, '00000000-0000-0000-0000-000000000000'::uuid),
|
||||
NULLIF(u.disconnect_reason, ''),
|
||||
NULLIF(u.disconnect_time, '0001-01-01 00:00:00Z'::timestamptz)
|
||||
FROM (
|
||||
SELECT
|
||||
unnest($1::uuid[]) AS id,
|
||||
unnest($2::timestamptz[]) AS connect_time,
|
||||
unnest($3::uuid[]) AS organization_id,
|
||||
unnest($4::uuid[]) AS workspace_owner_id,
|
||||
unnest($5::uuid[]) AS workspace_id,
|
||||
unnest($6::text[]) AS workspace_name,
|
||||
unnest($7::text[]) AS agent_name,
|
||||
unnest($8::connection_type[]) AS type,
|
||||
unnest($9::int4[]) AS code,
|
||||
unnest($10::bool[]) AS code_valid,
|
||||
unnest($11::inet[]) AS ip,
|
||||
unnest($12::text[]) AS user_agent,
|
||||
unnest($13::uuid[]) AS user_id,
|
||||
unnest($14::text[]) AS slug_or_port,
|
||||
unnest($15::uuid[]) AS connection_id,
|
||||
unnest($16::text[]) AS disconnect_reason,
|
||||
unnest($17::timestamptz[]) AS disconnect_time
|
||||
) AS u
|
||||
ON CONFLICT (connection_id, workspace_id, agent_name)
|
||||
DO UPDATE SET
|
||||
-- Pick the earliest real connect_time. The zero sentinel
|
||||
-- ('0001-01-01') means the batch didn't know the connect_time
|
||||
-- (e.g. a pure disconnect event), so we keep the existing value.
|
||||
connect_time = CASE
|
||||
WHEN EXCLUDED.connect_time = '0001-01-01 00:00:00Z'::timestamptz
|
||||
THEN connection_logs.connect_time
|
||||
WHEN connection_logs.connect_time = '0001-01-01 00:00:00Z'::timestamptz
|
||||
THEN EXCLUDED.connect_time
|
||||
ELSE LEAST(connection_logs.connect_time, EXCLUDED.connect_time)
|
||||
END,
|
||||
disconnect_time = CASE
|
||||
WHEN connection_logs.disconnect_time IS NULL
|
||||
THEN EXCLUDED.disconnect_time
|
||||
ELSE connection_logs.disconnect_time
|
||||
END,
|
||||
disconnect_reason = CASE
|
||||
WHEN connection_logs.disconnect_reason IS NULL
|
||||
THEN EXCLUDED.disconnect_reason
|
||||
ELSE connection_logs.disconnect_reason
|
||||
END,
|
||||
code = CASE
|
||||
WHEN connection_logs.code IS NULL
|
||||
THEN EXCLUDED.code
|
||||
ELSE connection_logs.code
|
||||
END
|
||||
`
|
||||
|
||||
type BatchUpsertConnectionLogsParams struct {
|
||||
ID []uuid.UUID `db:"id" json:"id"`
|
||||
ConnectTime []time.Time `db:"connect_time" json:"connect_time"`
|
||||
OrganizationID []uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
WorkspaceOwnerID []uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
|
||||
WorkspaceID []uuid.UUID `db:"workspace_id" json:"workspace_id"`
|
||||
WorkspaceName []string `db:"workspace_name" json:"workspace_name"`
|
||||
AgentName []string `db:"agent_name" json:"agent_name"`
|
||||
Type []ConnectionType `db:"type" json:"type"`
|
||||
Code []int32 `db:"code" json:"code"`
|
||||
CodeValid []bool `db:"code_valid" json:"code_valid"`
|
||||
Ip []pqtype.Inet `db:"ip" json:"ip"`
|
||||
UserAgent []string `db:"user_agent" json:"user_agent"`
|
||||
UserID []uuid.UUID `db:"user_id" json:"user_id"`
|
||||
SlugOrPort []string `db:"slug_or_port" json:"slug_or_port"`
|
||||
ConnectionID []uuid.UUID `db:"connection_id" json:"connection_id"`
|
||||
DisconnectReason []string `db:"disconnect_reason" json:"disconnect_reason"`
|
||||
DisconnectTime []time.Time `db:"disconnect_time" json:"disconnect_time"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) BatchUpsertConnectionLogs(ctx context.Context, arg BatchUpsertConnectionLogsParams) error {
|
||||
_, err := q.db.ExecContext(ctx, batchUpsertConnectionLogs,
|
||||
pq.Array(arg.ID),
|
||||
pq.Array(arg.ConnectTime),
|
||||
pq.Array(arg.OrganizationID),
|
||||
pq.Array(arg.WorkspaceOwnerID),
|
||||
pq.Array(arg.WorkspaceID),
|
||||
pq.Array(arg.WorkspaceName),
|
||||
pq.Array(arg.AgentName),
|
||||
pq.Array(arg.Type),
|
||||
pq.Array(arg.Code),
|
||||
pq.Array(arg.CodeValid),
|
||||
pq.Array(arg.Ip),
|
||||
pq.Array(arg.UserAgent),
|
||||
pq.Array(arg.UserID),
|
||||
pq.Array(arg.SlugOrPort),
|
||||
pq.Array(arg.ConnectionID),
|
||||
pq.Array(arg.DisconnectReason),
|
||||
pq.Array(arg.DisconnectTime),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
const countConnectionLogs = `-- name: CountConnectionLogs :one
|
||||
SELECT
|
||||
COUNT(*) AS count
|
||||
@@ -7753,120 +7870,6 @@ func (q *sqlQuerier) GetConnectionLogsOffset(ctx context.Context, arg GetConnect
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const upsertConnectionLog = `-- name: UpsertConnectionLog :one
|
||||
INSERT INTO connection_logs (
|
||||
id,
|
||||
connect_time,
|
||||
organization_id,
|
||||
workspace_owner_id,
|
||||
workspace_id,
|
||||
workspace_name,
|
||||
agent_name,
|
||||
type,
|
||||
code,
|
||||
ip,
|
||||
user_agent,
|
||||
user_id,
|
||||
slug_or_port,
|
||||
connection_id,
|
||||
disconnect_reason,
|
||||
disconnect_time
|
||||
) VALUES
|
||||
($1, $15, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14,
|
||||
-- If we've only received a disconnect event, mark the event as immediately
|
||||
-- closed.
|
||||
CASE
|
||||
WHEN $16::connection_status = 'disconnected'
|
||||
THEN $15 :: timestamp with time zone
|
||||
ELSE NULL
|
||||
END)
|
||||
ON CONFLICT (connection_id, workspace_id, agent_name)
|
||||
DO UPDATE SET
|
||||
-- No-op if the connection is still open.
|
||||
disconnect_time = CASE
|
||||
WHEN $16::connection_status = 'disconnected'
|
||||
-- Can only be set once
|
||||
AND connection_logs.disconnect_time IS NULL
|
||||
THEN EXCLUDED.connect_time
|
||||
ELSE connection_logs.disconnect_time
|
||||
END,
|
||||
disconnect_reason = CASE
|
||||
WHEN $16::connection_status = 'disconnected'
|
||||
-- Can only be set once
|
||||
AND connection_logs.disconnect_reason IS NULL
|
||||
THEN EXCLUDED.disconnect_reason
|
||||
ELSE connection_logs.disconnect_reason
|
||||
END,
|
||||
code = CASE
|
||||
WHEN $16::connection_status = 'disconnected'
|
||||
-- Can only be set once
|
||||
AND connection_logs.code IS NULL
|
||||
THEN EXCLUDED.code
|
||||
ELSE connection_logs.code
|
||||
END
|
||||
RETURNING id, connect_time, organization_id, workspace_owner_id, workspace_id, workspace_name, agent_name, type, ip, code, user_agent, user_id, slug_or_port, connection_id, disconnect_time, disconnect_reason
|
||||
`
|
||||
|
||||
type UpsertConnectionLogParams struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
|
||||
WorkspaceOwnerID uuid.UUID `db:"workspace_owner_id" json:"workspace_owner_id"`
|
||||
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
|
||||
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
|
||||
AgentName string `db:"agent_name" json:"agent_name"`
|
||||
Type ConnectionType `db:"type" json:"type"`
|
||||
Code sql.NullInt32 `db:"code" json:"code"`
|
||||
Ip pqtype.Inet `db:"ip" json:"ip"`
|
||||
UserAgent sql.NullString `db:"user_agent" json:"user_agent"`
|
||||
UserID uuid.NullUUID `db:"user_id" json:"user_id"`
|
||||
SlugOrPort sql.NullString `db:"slug_or_port" json:"slug_or_port"`
|
||||
ConnectionID uuid.NullUUID `db:"connection_id" json:"connection_id"`
|
||||
DisconnectReason sql.NullString `db:"disconnect_reason" json:"disconnect_reason"`
|
||||
Time time.Time `db:"time" json:"time"`
|
||||
ConnectionStatus ConnectionStatus `db:"connection_status" json:"connection_status"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertConnectionLog(ctx context.Context, arg UpsertConnectionLogParams) (ConnectionLog, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertConnectionLog,
|
||||
arg.ID,
|
||||
arg.OrganizationID,
|
||||
arg.WorkspaceOwnerID,
|
||||
arg.WorkspaceID,
|
||||
arg.WorkspaceName,
|
||||
arg.AgentName,
|
||||
arg.Type,
|
||||
arg.Code,
|
||||
arg.Ip,
|
||||
arg.UserAgent,
|
||||
arg.UserID,
|
||||
arg.SlugOrPort,
|
||||
arg.ConnectionID,
|
||||
arg.DisconnectReason,
|
||||
arg.Time,
|
||||
arg.ConnectionStatus,
|
||||
)
|
||||
var i ConnectionLog
|
||||
err := row.Scan(
|
||||
&i.ID,
|
||||
&i.ConnectTime,
|
||||
&i.OrganizationID,
|
||||
&i.WorkspaceOwnerID,
|
||||
&i.WorkspaceID,
|
||||
&i.WorkspaceName,
|
||||
&i.AgentName,
|
||||
&i.Type,
|
||||
&i.Ip,
|
||||
&i.Code,
|
||||
&i.UserAgent,
|
||||
&i.UserID,
|
||||
&i.SlugOrPort,
|
||||
&i.ConnectionID,
|
||||
&i.DisconnectTime,
|
||||
&i.DisconnectReason,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const deleteCryptoKey = `-- name: DeleteCryptoKey :one
|
||||
UPDATE crypto_keys
|
||||
SET secret = NULL, secret_key_id = NULL
|
||||
|
||||
Reference in New Issue
Block a user