diff --git a/api/api.go b/api/api.go index fef528b8..9ec72fa4 100644 --- a/api/api.go +++ b/api/api.go @@ -51,6 +51,7 @@ func Routes() []*router.Route { router.NewRoute("POST", "/auth/upgrade-guest-existing", handlers.UpgradeGuestExisting), router.NewRoute("POST", "/network/auth-client", handlers.AuthNetworkClient), router.NewRoute("POST", "/network/remove-client", handlers.RemoveNetworkClient), + router.NewRoute("POST", "/network/remove-clients", handlers.RemoveNetworkClients), router.NewRoute("GET", "/network/clients", handlers.NetworkClients), router.NewRoute("GET", "/network/peers", handlers.NetworkPeers), router.NewRoute("GET", "/network/provider-locations", handlers.NetworkGetProviderLocations), diff --git a/api/handlers/network_client_handlers.go b/api/handlers/network_client_handlers.go index 59ae5b9c..50c558ff 100644 --- a/api/handlers/network_client_handlers.go +++ b/api/handlers/network_client_handlers.go @@ -38,6 +38,10 @@ func RemoveNetworkClient(w http.ResponseWriter, r *http.Request) { router.WrapWithInputRequireAuth(model.RemoveNetworkClient, w, r) } +func RemoveNetworkClients(w http.ResponseWriter, r *http.Request) { + router.WrapWithInputRequireAuth(model.RemoveNetworkClients, w, r) +} + func RemoveNetwork(w http.ResponseWriter, r *http.Request) { router.WrapRequireAuth(controller.NetworkRemove, w, r) } diff --git a/model/network_client_model.go b/model/network_client_model.go index 0bb6a98a..9eb5d27a 100644 --- a/model/network_client_model.go +++ b/model/network_client_model.go @@ -27,6 +27,7 @@ import ( "github.com/urnetwork/server" "github.com/urnetwork/server/session" + "github.com/urnetwork/server/task" // "github.com/urnetwork/server/ulid" // "github.com/urnetwork/server/jwt" "github.com/urnetwork/connect" @@ -699,6 +700,250 @@ func RemoveNetworkClient( return removeClientResult, removeClientErr } +// matches the batch size `RemoveDisconnectedNetworkClients` already uses for +// bounded maintenance sweeps of this same table (see `markTopLevelBatchCount` +// above): large enough to make a real dent per transaction, small enough that +// no single transaction runs long or holds locks for long. +const RemoveNetworkClientsBatchCount = 10000 + +func removeNetworkClientsBatchExec(ctx context.Context, tx server.PgTx, clientIds []server.Id, networkId server.Id) { + _, err := tx.Exec( + ctx, + ` + UPDATE network_client + SET + active = false, + deactivate_time = $3 + WHERE + client_id = ANY($1) AND + network_id = $2 + `, + clientIds, + networkId, + server.NowUtc(), + ) + server.Raise(err) +} + +type RemoveNetworkClientsBatchArgs struct { + ClientIds []server.Id `json:"client_ids"` +} + +type RemoveNetworkClientsBatchResult struct{} + +// RemoveNetworkClientsBatch deactivates up to RemoveNetworkClientsBatchCount +// clients in a single transaction on the regular request pool. Callers with +// more ids than that should go through RemoveNetworkClients, which routes +// large requests to the background task instead of calling this directly. +// TxReadCommitted matches RemoveDisconnectedNetworkClients's isolation level +// for this same table, avoiding serialization-failure panics under +// concurrent overlapping updates that the pool's default isolation would risk. +func RemoveNetworkClientsBatch( + removeClients *RemoveNetworkClientsBatchArgs, + session *session.ClientSession, +) (*RemoveNetworkClientsBatchResult, error) { + server.Tx(session.Ctx, func(tx server.PgTx) { + removeNetworkClientsBatchExec(session.Ctx, tx, removeClients.ClientIds, session.ByJwt.NetworkId) + }, server.TxReadCommitted) + return &RemoveNetworkClientsBatchResult{}, nil +} + +// outer sanity bound on a single request's payload, not a realistic +// operating limit (the known real-world case is ~400k-1M): guards against an +// unbounded/malformed request forcing an arbitrarily large synchronous +// json.Marshal + single-row task args_json write in RemoveNetworkClients, and +// against an unbounded in-memory slice being held for the life of a task. +// Kept at 2x the known real-world ceiling rather than higher: at ~38 bytes +// per id as JSON, this is already ~75MB in the worst case (a single +// pending_task.args_json write on the request path), so the bound is sized +// to limit that cost, not just to be a round number. +const MaxRemoveNetworkClientsCount = 2000000 + +type RemoveNetworkClientsArgs struct { + ClientIds []server.Id `json:"client_ids"` +} + +type RemoveNetworkClientsResult struct { + // true if the request was handed off to the background task instead of + // being applied synchronously; deactivation is not yet guaranteed + // complete when this is true + Scheduled bool `json:"scheduled,omitempty"` + // true if a background bulk-delete run for this network was already in + // progress, so this request's ids were NOT scheduled; the caller should + // wait for the in-progress run to finish and retry + AlreadyInProgress bool `json:"already_in_progress,omitempty"` +} + +// runNetworkClientsTaskKey is the run_once key scoping "one background +// bulk-delete run per network at a time". It's shared between the initial +// schedule in RemoveNetworkClients and the continuation reschedule in +// RemoveNetworkClientsTaskPost, so the key stays held for the full duration +// of a (possibly multi-invocation) run, not just its first invocation. +func runNetworkClientsTaskKey(networkId server.Id) *task.RunOnceOption { + return task.RunOnce("model.RemoveNetworkClientsTask", networkId) +} + +// RemoveNetworkClients deactivates the given clients, scoped to the caller's +// network. Small requests are applied synchronously in one transaction. +// Requests larger than RemoveNetworkClientsBatchCount are handed off to +// RemoveNetworkClientsTask, a background task that processes the full list in +// bounded batches, so a single call can clear a network with hundreds of +// thousands (or millions) of offline clients without holding a long-running +// transaction on the request path or forcing the caller to chunk the request +// themselves. +func RemoveNetworkClients( + removeClients *RemoveNetworkClientsArgs, + session *session.ClientSession, +) (*RemoveNetworkClientsResult, error) { + if len(removeClients.ClientIds) == 0 { + return &RemoveNetworkClientsResult{}, nil + } + if MaxRemoveNetworkClientsCount < len(removeClients.ClientIds) { + return nil, fmt.Errorf("Too many client ids (max %d).", MaxRemoveNetworkClientsCount) + } + + if len(removeClients.ClientIds) <= RemoveNetworkClientsBatchCount { + _, err := RemoveNetworkClientsBatch(&RemoveNetworkClientsBatchArgs{ + ClientIds: removeClients.ClientIds, + }, session) + if err != nil { + return nil, err + } + return &RemoveNetworkClientsResult{}, nil + } + + // one background bulk-delete run per network at a time. + // ScheduleTaskInTxIfAbsent (unlike plain ScheduleTask+RunOnce) makes the + // "only if not already pending" check atomic with the insert -- a single + // `INSERT ... ON CONFLICT (run_once_key) DO NOTHING`, reporting whether + // the row was actually inserted. This closes the race a naive + // check-then-act would have: with check-then-act, two near-simultaneous + // requests for the same network could both pass the check, and the + // second's schedule call would then hit ON CONFLICT DO UPDATE -- which + // merges only timing/priority into the existing row, NOT args_json -- + // silently dropping the second call's client_ids while still reporting + // success. The atomic insert-or-detect-conflict here means a duplicate + // is always rejected outright, never silently swallowed. + scheduled, _ := task.ScheduleTaskIfAbsent( + RemoveNetworkClientsTask, + &RemoveNetworkClientsTaskArgs{ + ClientIds: removeClients.ClientIds, + }, + session, + runNetworkClientsTaskKey(session.ByJwt.NetworkId), + // bulk cleanup must never compete with revenue/critical-path tasks + // (payouts, contract close) under multi-tenant load + task.Priority(task.TaskPrioritySlowest), + task.MaxTime(30*time.Minute), + ) + if !scheduled { + return &RemoveNetworkClientsResult{AlreadyInProgress: true}, nil + } + + return &RemoveNetworkClientsResult{Scheduled: true}, nil +} + +// number of batches processed per RemoveNetworkClientsTask invocation before +// it self-reschedules the remainder (see RemoveNetworkClientsTaskPost), +// instead of looping over the entire list in one invocation. This bounds +// three things: (1) a single invocation's duration stays well inside +// MaxTime regardless of total list size, (2) if an invocation is retried +// after a crash/timeout, it only re-attempts this bounded chunk, not the +// entire original list, so a run makes durable checkpointed progress instead +// of restarting from the beginning every retry, and (3) a persistently +// failing chunk only blocks its own continuation, not the network's ability +// to ever complete a bulk delete. +const RemoveNetworkClientsTaskBatchLimit = 20 + +type RemoveNetworkClientsTaskArgs struct { + ClientIds []server.Id `json:"client_ids"` +} + +type RemoveNetworkClientsTaskResult struct { + // ids not yet processed by this invocation; if non-empty, + // RemoveNetworkClientsTaskPost reschedules them as a new task under the + // same run_once key + RemainingClientIds []server.Id `json:"remaining_client_ids,omitempty"` +} + +// RemoveNetworkClientsTask is the background counterpart to +// RemoveNetworkClients for requests larger than RemoveNetworkClientsBatchCount. +// It runs on the taskworker service (not the live request pool) and processes +// up to RemoveNetworkClientsTaskBatchLimit batches of +// RemoveNetworkClientsBatchCount ids, each its own short server.MaintenanceTx +// transaction at TxReadCommitted, mirroring the batching and isolation level +// RemoveDisconnectedNetworkClients already uses for this table. The update is +// idempotent, so a retried or re-run invocation is safe. +func RemoveNetworkClientsTask( + removeClients *RemoveNetworkClientsTaskArgs, + session *session.ClientSession, +) (*RemoveNetworkClientsTaskResult, error) { + if session.ByJwt == nil { + // unreachable in the normal flow (always scheduled from an + // authenticated handler with ByJwt set), but this task can outlive + // the request that scheduled it -- fail loudly rather than panic on + // a nil dereference if a pending_task row is ever reconstructed + // with an empty/corrupted client_by_jwt_json + return nil, fmt.Errorf("Missing network for bulk client removal.") + } + + clientIds := removeClients.ClientIds + batches := 0 + for 0 < len(clientIds) && batches < RemoveNetworkClientsTaskBatchLimit { + batchCount := RemoveNetworkClientsBatchCount + if len(clientIds) < batchCount { + batchCount = len(clientIds) + } + batch := clientIds[:batchCount] + clientIds = clientIds[batchCount:] + + server.MaintenanceTx(session.Ctx, func(tx server.PgTx) { + removeNetworkClientsBatchExec(session.Ctx, tx, batch, session.ByJwt.NetworkId) + }, server.TxReadCommitted) + batches += 1 + } + + return &RemoveNetworkClientsTaskResult{ + RemainingClientIds: clientIds, + }, nil +} + +// RemoveNetworkClientsTaskPost reschedules any remainder left by +// RemoveNetworkClientsTask under the SAME run_once key as the original +// request, so the "one run per network" guarantee holds for the full +// duration of a multi-invocation run and not just its first invocation -- +// otherwise the key would free up as soon as the first chunk finished, even +// though most of the list might still be unprocessed, and a second +// concurrent request would incorrectly be allowed through. +// +// This reschedules with plain ScheduleTaskInTx (ON CONFLICT DO UPDATE), not +// ScheduleTaskInTxIfAbsent (ON CONFLICT DO NOTHING): the finishing task's own +// pending_task row is deleted as part of this same evaluation transaction, +// so the run_once key is guaranteed free at the point this INSERT runs -- +// there's no possible conflict with a legitimate concurrent request here, +// only with the row this exact chain just vacated. +func RemoveNetworkClientsTaskPost( + removeClients *RemoveNetworkClientsTaskArgs, + result *RemoveNetworkClientsTaskResult, + session *session.ClientSession, + tx server.PgTx, +) error { + if 0 < len(result.RemainingClientIds) { + task.ScheduleTaskInTx( + tx, + RemoveNetworkClientsTask, + &RemoveNetworkClientsTaskArgs{ + ClientIds: result.RemainingClientIds, + }, + session, + runNetworkClientsTaskKey(session.ByJwt.NetworkId), + task.Priority(task.TaskPrioritySlowest), + task.MaxTime(30*time.Minute), + ) + } + return nil +} + type NetworkClientsResult struct { Clients []*NetworkClientInfo `json:"clients"` } diff --git a/model/network_client_model_test.go b/model/network_client_model_test.go index 81826e84..0a93603a 100644 --- a/model/network_client_model_test.go +++ b/model/network_client_model_test.go @@ -12,6 +12,7 @@ import ( "github.com/urnetwork/server" "github.com/urnetwork/server/jwt" "github.com/urnetwork/server/session" + "github.com/urnetwork/server/task" ) func TestNetworkClientHandlerLifecycle(t *testing.T) { @@ -200,6 +201,733 @@ func TestSetProvide(t *testing.T) { }) } +// RemoveNetworkClients must accept a uuid[]-bound array of ids that don't +// exist yet (nothing to match), and must not error or panic on the +// []server.Id -> uuid[] cast. +func TestRemoveNetworkClientsUUIDArrayBinding(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + // random ids with no matching rows + clientIds := []server.Id{server.NewId(), server.NewId(), server.NewId()} + + args := &RemoveNetworkClientsArgs{ + ClientIds: clientIds, + } + + result, err := RemoveNetworkClients(args, sess) + assert.Equal(t, err, nil) + assert.NotEqual(t, result, nil) + }) +} + +// The empty-ids case must be a no-op that returns cleanly without opening a +// transaction (an empty `ANY($1)` array matches nothing, but there's no +// reason to pay for the round trip). +func TestRemoveNetworkClientsEmptyIdsNoop(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + deviceId := server.NewId() + clientId := server.NewId() + + Testing_CreateDevice(ctx, networkId, deviceId, clientId, "test", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: []server.Id{}, + }, sess) + assert.Equal(t, err, nil) + assert.NotEqual(t, result, nil) + + // the existing client must be untouched + assert.NotEqual(t, GetNetworkClient(ctx, clientId), nil) + }) +} + +// The bulk path must actually deactivate the targeted clients (mirroring the +// single-client RemoveNetworkClient behavior) while leaving other clients on +// the same network alone. +func TestRemoveNetworkClientsDeactivatesTargetedClients(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + deviceIdA := server.NewId() + clientIdA := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceIdA, clientIdA, "test-a", "test") + + deviceIdB := server.NewId() + clientIdB := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceIdB, clientIdB, "test-b", "test") + + deviceIdC := server.NewId() + clientIdC := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceIdC, clientIdC, "test-c", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + beforeCall := server.NowUtc() + _, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: []server.Id{clientIdA, clientIdB}, + }, sess) + assert.Equal(t, err, nil) + + // targeted clients are deactivated + assert.Equal(t, GetNetworkClient(ctx, clientIdA), nil) + assert.Equal(t, GetNetworkClient(ctx, clientIdB), nil) + + // the untargeted client on the same network is untouched + assert.NotEqual(t, GetNetworkClient(ctx, clientIdC), nil) + + // deactivate_time must be stamped, matching the single-client path, + // so the reap job (`COALESCE(deactivate_time, create_time)`) applies + // the grace period from the actual deactivation instead of falling + // back to create_time + var deactivateTime *time.Time + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + `SELECT deactivate_time FROM network_client WHERE client_id = $1`, + clientIdA, + ) + server.WithPgResult(result, err, func() { + if result.Next() { + server.Raise(result.Scan(&deactivateTime)) + } + }) + }) + if deactivateTime == nil { + t.Fatal("deactivate_time was not set") + } + if deactivateTime.Before(beforeCall) { + t.Fatal("deactivate_time predates the call") + } + }) +} + +// A request at or under the batch size must be applied synchronously: the +// caller gets a definite "done" (Scheduled == false) and the clients are +// already deactivated by the time the call returns. +func TestRemoveNetworkClientsSmallRequestIsSynchronous(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + deviceId := server.NewId() + clientId := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceId, clientId, "test", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: []server.Id{clientId}, + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, result.Scheduled, false) + + // already deactivated, not just enqueued + assert.Equal(t, GetNetworkClient(ctx, clientId), nil) + }) +} + +// A request over the batch size must be handed off to the background task +// instead of run inline: the caller gets Scheduled == true and the clients +// are NOT yet deactivated when the call returns (only the task is enqueued; +// RemoveNetworkClientsTask is what actually applies it, tested separately +// below). +func TestRemoveNetworkClientsLargeRequestIsScheduled(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + deviceId := server.NewId() + clientId := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceId, clientId, "test", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + clientIds := make([]server.Id, RemoveNetworkClientsBatchCount+1) + clientIds[0] = clientId + for i := 1; i < len(clientIds); i++ { + clientIds[i] = server.NewId() + } + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds, + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, result.Scheduled, true) + assert.Equal(t, result.AlreadyInProgress, false) + + // not yet applied: this call only enqueued the background task + assert.NotEqual(t, GetNetworkClient(ctx, clientId), nil) + }) +} + +// A second large request for the same network, made while the first is +// still pending, must be rejected outright (AlreadyInProgress == true, not +// Scheduled) rather than silently merged into the first task -- merging +// would drop the second call's client_ids, since ScheduleTask's run_once +// conflict path only updates timing/priority, not args. +func TestRemoveNetworkClientsRejectsDuplicateInProgress(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + firstDeviceId := server.NewId() + firstClientId := server.NewId() + Testing_CreateDevice(ctx, networkId, firstDeviceId, firstClientId, "first", "test") + + secondDeviceId := server.NewId() + secondClientId := server.NewId() + Testing_CreateDevice(ctx, networkId, secondDeviceId, secondClientId, "second", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + largeClientIds := func(first server.Id) []server.Id { + clientIds := make([]server.Id, RemoveNetworkClientsBatchCount+1) + clientIds[0] = first + for i := 1; i < len(clientIds); i++ { + clientIds[i] = server.NewId() + } + return clientIds + } + + firstResult, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: largeClientIds(firstClientId), + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, firstResult.Scheduled, true) + assert.Equal(t, firstResult.AlreadyInProgress, false) + + // a second large request for the same network, while the first is + // still pending (unclaimed), must be rejected, not scheduled + secondResult, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: largeClientIds(secondClientId), + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, secondResult.Scheduled, false) + assert.Equal(t, secondResult.AlreadyInProgress, true) + + // neither client is deactivated yet: the first task hasn't run + // (only enqueued), and the second call was never scheduled at all + assert.NotEqual(t, GetNetworkClient(ctx, firstClientId), nil) + assert.NotEqual(t, GetNetworkClient(ctx, secondClientId), nil) + }) +} + +// A large request for a DIFFERENT network must not be blocked by an +// in-progress run on another network: the run_once key is scoped per network. +func TestRemoveNetworkClientsInProgressIsScopedPerNetwork(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkIdA := server.NewId() + networkIdB := server.NewId() + + deviceIdA := server.NewId() + clientIdA := server.NewId() + Testing_CreateDevice(ctx, networkIdA, deviceIdA, clientIdA, "a", "test") + + deviceIdB := server.NewId() + clientIdB := server.NewId() + Testing_CreateDevice(ctx, networkIdB, deviceIdB, clientIdB, "b", "test") + + largeClientIds := func(first server.Id) []server.Id { + clientIds := make([]server.Id, RemoveNetworkClientsBatchCount+1) + clientIds[0] = first + for i := 1; i < len(clientIds); i++ { + clientIds[i] = server.NewId() + } + return clientIds + } + + sessA := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{NetworkId: networkIdA}, + } + sessB := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{NetworkId: networkIdB}, + } + + resultA, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: largeClientIds(clientIdA), + }, sessA) + assert.Equal(t, err, nil) + assert.Equal(t, resultA.Scheduled, true) + + resultB, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: largeClientIds(clientIdB), + }, sessB) + assert.Equal(t, err, nil) + assert.Equal(t, resultB.Scheduled, true) + assert.Equal(t, resultB.AlreadyInProgress, false) + }) +} + +// A request over the outer sanity cap must be rejected outright, before any +// scheduling or database work -- this is a payload-size guard, not a +// realistic operating limit (known real-world usage is far below it). +func TestRemoveNetworkClientsRejectsOversizedRequest(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + // zero-value ids are fine here: the request must be rejected on + // length alone, before any id is ever read + clientIds := make([]server.Id, MaxRemoveNetworkClientsCount+1) + + _, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds, + }, sess) + assert.NotEqual(t, err, nil) + }) +} + +// End-to-end through a real task.TaskWorker (not calling RemoveNetworkClientsTask +// directly): a request large enough to require more than one invocation +// (RemoveNetworkClientsTaskBatchLimit batches) must fully deactivate every +// client across the self-rescheduled chain; a duplicate large request for the +// same network must be rejected while any part of the chain is still +// in-flight; and once the whole chain completes, the run_once key must be +// free again for a new request. This is the scenario the AlreadyInProgress +// flag depends on, and it isn't exercised by calling RemoveNetworkClientsTask +// directly (that only tests one invocation's batching, not the pending_task +// lifecycle across a real claim/run/post/reschedule cycle). +func TestRemoveNetworkClientsTaskLifecycleThroughRealWorker(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + // exceed one invocation's batch limit so the task must self-reschedule + // at least once + targetCount := RemoveNetworkClientsTaskBatchLimit*RemoveNetworkClientsBatchCount + 5000 + var clientIds []server.Id + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + ` + INSERT INTO network_client (client_id, network_id, active, create_time, auth_time) + SELECT gen_random_uuid(), $1, true, now(), now() + FROM generate_series(1, $2) g + RETURNING client_id + `, + networkId, + targetCount, + ) + server.WithPgResult(result, err, func() { + for result.Next() { + var clientId server.Id + server.Raise(result.Scan(&clientId)) + clientIds = append(clientIds, clientId) + } + }) + }) + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds, + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, result.Scheduled, true) + + // while the run is in flight, a second large request for the same + // network must be rejected, not merged or double-scheduled + dupResult, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds[:RemoveNetworkClientsBatchCount+1], + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, dupResult.Scheduled, false) + assert.Equal(t, dupResult.AlreadyInProgress, true) + + taskWorker := task.NewTaskWorkerWithDefaults(ctx) + taskWorker.AddTargets(task.NewTaskTargetWithPost(RemoveNetworkClientsTask, RemoveNetworkClientsTaskPost)) + + // drive the worker until no pending work remains for this run; bounded + // so a stuck run fails the test instead of hanging it + for i := 0; i < 50; i++ { + finishedTaskIds, rescheduledTaskIds, postRescheduledTaskIds, err := taskWorker.EvalTasks(10) + assert.Equal(t, err, nil) + if len(finishedTaskIds)+len(rescheduledTaskIds)+len(postRescheduledTaskIds) == 0 && + len(task.ListPendingTasks(ctx)) == 0 { + break + } + } + if 0 < len(task.ListPendingTasks(ctx)) { + t.Fatal("run did not complete within the bounded number of eval passes") + } + + var remainingActive int + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + `SELECT COUNT(*) FROM network_client WHERE network_id = $1 AND active = true`, + networkId, + ) + server.WithPgResult(result, err, func() { + if result.Next() { + server.Raise(result.Scan(&remainingActive)) + } + }) + }) + assert.Equal(t, remainingActive, 0) + + // deactivate_time must be stamped across the whole async chain, not + // just the first invocation -- this is the only test exercising + // multiple RemoveNetworkClientsTask invocations, so it's the one + // that would catch a continuation batch losing the stamp + var missingDeactivateTime int + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + `SELECT COUNT(*) FROM network_client WHERE network_id = $1 AND active = false AND deactivate_time IS NULL`, + networkId, + ) + server.WithPgResult(result, err, func() { + if result.Next() { + server.Raise(result.Scan(&missingDeactivateTime)) + } + }) + }) + assert.Equal(t, missingDeactivateTime, 0) + + // the run_once key must be free again now that the full chain + // finished, so a new large request for this network succeeds + afterResult, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds[:RemoveNetworkClientsBatchCount+1], + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, afterResult.Scheduled, true) + }) +} + +// RemoveNetworkClientsTask (the background counterpart RemoveNetworkClients +// hands large requests off to) must deactivate every targeted client across +// multiple internal batches, and must not touch a client on a different +// network that happens to fall in the same id range. Called directly here +// (not through the task queue) to test its batching logic deterministically +// without a running taskworker. +func TestRemoveNetworkClientsTaskSpansMultipleBatches(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + otherNetworkId := server.NewId() + + // seed more than 2x the batch size for the target network, plus one + // client on a different network, via a bulk insert (matches the + // generate_series seeding pattern used elsewhere in this package) + targetCount := 2*RemoveNetworkClientsBatchCount + 5000 + var clientIds []server.Id + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + ` + INSERT INTO network_client (client_id, network_id, active, create_time, auth_time) + SELECT gen_random_uuid(), $1, true, now(), now() + FROM generate_series(1, $2) g + RETURNING client_id + `, + networkId, + targetCount, + ) + server.WithPgResult(result, err, func() { + for result.Next() { + var clientId server.Id + server.Raise(result.Scan(&clientId)) + clientIds = append(clientIds, clientId) + } + }) + }) + + otherDeviceId := server.NewId() + otherClientId := server.NewId() + Testing_CreateDevice(ctx, otherNetworkId, otherDeviceId, otherClientId, "other", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + _, err := RemoveNetworkClientsTask(&RemoveNetworkClientsTaskArgs{ + ClientIds: clientIds, + }, sess) + assert.Equal(t, err, nil) + + var remainingActive int + server.Db(ctx, func(conn server.PgConn) { + result, err := conn.Query( + ctx, + `SELECT COUNT(*) FROM network_client WHERE network_id = $1 AND active = true`, + networkId, + ) + server.WithPgResult(result, err, func() { + if result.Next() { + server.Raise(result.Scan(&remainingActive)) + } + }) + }) + assert.Equal(t, remainingActive, 0) + + // a client on a different network is untouched + assert.NotEqual(t, GetNetworkClient(ctx, otherClientId), nil) + }) +} + +// If a task worker retries RemoveNetworkClientsTask (e.g. after a crash or a +// MaxTime timeout mid-run), re-running it with the same args must be a safe +// no-op the second time: `active = false` on an already-inactive row doesn't +// error, and deactivate_time just advances rather than corrupting state. +func TestRemoveNetworkClientsTaskIsIdempotentOnRetry(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + deviceId := server.NewId() + clientId := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceId, clientId, "test", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + args := &RemoveNetworkClientsTaskArgs{ + ClientIds: []server.Id{clientId}, + } + + _, err := RemoveNetworkClientsTask(args, sess) + assert.Equal(t, err, nil) + assert.Equal(t, GetNetworkClient(ctx, clientId), nil) + + // simulate a retry of the same task with the same args + _, err = RemoveNetworkClientsTask(args, sess) + assert.Equal(t, err, nil) + assert.Equal(t, GetNetworkClient(ctx, clientId), nil) + }) +} + +// RemoveNetworkClientsTask must fail loudly (not panic on a nil dereference +// inside a MaintenanceTx) if it's ever run with a session that has no +// ByJwt -- unreachable in the normal flow (always scheduled from an +// authenticated handler), but the task can outlive the request that +// scheduled it, so this guards against a corrupted/reconstructed session. +func TestRemoveNetworkClientsTaskRejectsNilByJwt(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: nil, + } + + _, err := RemoveNetworkClientsTask(&RemoveNetworkClientsTaskArgs{ + ClientIds: []server.Id{server.NewId()}, + }, sess) + assert.NotEqual(t, err, nil) + }) +} + +// A request of exactly RemoveNetworkClientsBatchCount ids must still take +// the synchronous path (the boundary is "<=", not "<"). +func TestRemoveNetworkClientsExactlyAtBatchCountIsSynchronous(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + deviceId := server.NewId() + clientId := server.NewId() + Testing_CreateDevice(ctx, networkId, deviceId, clientId, "test", "test") + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + clientIds := make([]server.Id, RemoveNetworkClientsBatchCount) + clientIds[0] = clientId + for i := 1; i < len(clientIds); i++ { + clientIds[i] = server.NewId() + } + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds, + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, result.Scheduled, false) + + // already applied synchronously, not enqueued + assert.Equal(t, GetNetworkClient(ctx, clientId), nil) + }) +} + +// A request of exactly MaxRemoveNetworkClientsCount ids must be accepted +// (the boundary is "> max is rejected", not ">="). This only exercises +// scheduling (one INSERT), not actually running the task, since processing +// 2M ids isn't practical inside a unit test. +func TestRemoveNetworkClientsExactlyAtCapIsAccepted(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + // zero-value ids are fine: this only tests that scheduling itself is + // accepted at exactly the cap, not that the ids get processed + clientIds := make([]server.Id, MaxRemoveNetworkClientsCount) + + result, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: clientIds, + }, sess) + assert.Equal(t, err, nil) + assert.Equal(t, result.Scheduled, true) + }) +} + +// RemoveNetworkClientsTaskPost, tested directly and in isolation (not +// through a full task.TaskWorker run, which the slow real-worker lifecycle +// test above already covers end-to-end): a non-empty RemainingClientIds must +// result in exactly one new pending task, scheduled under the same run_once +// key as the original request; an empty RemainingClientIds must schedule +// nothing at all. +func TestRemoveNetworkClientsTaskPostReschedulesRemainder(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + networkId := server.NewId() + + sess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: networkId, + }, + } + + remainingClientIds := []server.Id{server.NewId(), server.NewId()} + + server.Tx(ctx, func(tx server.PgTx) { + err := RemoveNetworkClientsTaskPost( + &RemoveNetworkClientsTaskArgs{ClientIds: remainingClientIds}, + &RemoveNetworkClientsTaskResult{RemainingClientIds: remainingClientIds}, + sess, + tx, + ) + assert.Equal(t, err, nil) + }) + + // the reschedule must be blocked by the run_once key (same key the + // original request would have used), proving the continuation was + // scheduled under it + scheduledAgain, _ := task.ScheduleTaskIfAbsent( + RemoveNetworkClientsTask, + &RemoveNetworkClientsTaskArgs{ClientIds: []server.Id{server.NewId()}}, + sess, + runNetworkClientsTaskKey(networkId), + ) + assert.Equal(t, scheduledAgain, false) + + // a DIFFERENT network's key must be unaffected + otherNetworkId := server.NewId() + otherSess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{NetworkId: otherNetworkId}, + } + scheduledOther, _ := task.ScheduleTaskIfAbsent( + RemoveNetworkClientsTask, + &RemoveNetworkClientsTaskArgs{ClientIds: []server.Id{server.NewId()}}, + otherSess, + runNetworkClientsTaskKey(otherNetworkId), + ) + assert.Equal(t, scheduledOther, true) + }) +} + +// A caller must not be able to deactivate another network's clients by +// passing their ids in the request body: `network_id = $2` in the query +// must be enforced server-side from the session, not trusted from input. +func TestRemoveNetworkClientsEnforcesNetworkScoping(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + + victimNetworkId := server.NewId() + victimDeviceId := server.NewId() + victimClientId := server.NewId() + Testing_CreateDevice(ctx, victimNetworkId, victimDeviceId, victimClientId, "victim", "test") + + attackerNetworkId := server.NewId() + attackerSess := &session.ClientSession{ + Ctx: ctx, + ByJwt: &jwt.ByJwt{ + NetworkId: attackerNetworkId, + }, + } + + _, err := RemoveNetworkClients(&RemoveNetworkClientsArgs{ + ClientIds: []server.Id{victimClientId}, + }, attackerSess) + assert.Equal(t, err, nil) + + // the victim's client must still be active + assert.NotEqual(t, GetNetworkClient(ctx, victimClientId), nil) + }) +} + // GetProvideModes / GetProvideSecretKey must fall back to postgres when the // redis cache is cold (data written before the redis layer existed, or evicted). // Regression: a redis miss used to leak a non-nil error even though the db diff --git a/task/task.go b/task/task.go index 543be9e2..d73f5fbd 100644 --- a/task/task.go +++ b/task/task.go @@ -134,13 +134,23 @@ func ScheduleTask[T any, R any]( return } -func ScheduleTaskInTx[T any, R any]( - tx server.PgTx, +type preparedTask struct { + taskId server.Id + functionName string + argsJson []byte + byJwtJson *string + runAt time.Time + runOnceKey *string + priority TaskPriority + maxTimeSeconds int +} + +func prepareTask[T any, R any]( taskFunction TaskFunction[T, R], args T, clientSession *session.ClientSession, opts ...any, -) (taskId server.Id) { +) preparedTask { taskTarget := NewTaskTarget(taskFunction) argsJson, err := json.Marshal(args) @@ -195,11 +205,29 @@ func ScheduleTaskInTx[T any, R any]( runOnceKey_ := runOnce.String() runOnceKey = &runOnceKey_ } - maxTimeSeconds := int(runMaxTime.MaxTime / time.Second) - claimTime := time.Time{} + return preparedTask{ + taskId: server.NewId(), + functionName: taskTarget.TargetFunctionName(), + argsJson: argsJson, + byJwtJson: byJwtJson, + runAt: runAt.At.UTC(), + runOnceKey: runOnceKey, + priority: runPriority.Priority, + maxTimeSeconds: int(runMaxTime.MaxTime / time.Second), + } +} + +func ScheduleTaskInTx[T any, R any]( + tx server.PgTx, + taskFunction TaskFunction[T, R], + args T, + clientSession *session.ClientSession, + opts ...any, +) (taskId server.Id) { + p := prepareTask(taskFunction, args, clientSession, opts...) - taskId = server.NewId() + claimTime := time.Time{} server.RaisePgResult(tx.Exec( clientSession.Ctx, @@ -222,17 +250,91 @@ func ScheduleTaskInTx[T any, R any]( run_priority = LEAST(pending_task.run_priority, $8), run_max_time_seconds = GREATEST(pending_task.run_max_time_seconds, $9) `, - taskId, - taskTarget.TargetFunctionName(), - argsJson, + p.taskId, + p.functionName, + p.argsJson, clientSession.ClientAddress, - byJwtJson, - runAt.At.UTC(), - runOnceKey, - runPriority.Priority, - maxTimeSeconds, + p.byJwtJson, + p.runAt, + p.runOnceKey, + p.priority, + p.maxTimeSeconds, claimTime, )) + return p.taskId +} + +// ScheduleTaskInTxIfAbsent is like ScheduleTaskInTx but for callers that need +// an atomic "only schedule if not already pending under this key" guarantee, +// instead of RunOnce's merge-on-conflict semantics. RunOnce's +// `ON CONFLICT (run_once_key) DO UPDATE` only merges run_at/run_priority/ +// run_max_time_seconds into an existing pending row -- crucially not +// args_json -- so if two different calls share a run_once key while the +// first is still pending, scheduling both would silently drop the second +// call's args while still reporting success. This does a single +// `INSERT ... ON CONFLICT (run_once_key) DO NOTHING` and reports via +// `scheduled` whether the row was actually inserted, so the caller can +// reject a duplicate outright -- atomically, in one round trip -- instead of +// a separate check-then-act that can itself race. runOnce is required (not +// optional via opts) since the whole point is a key-scoped guarantee. +func ScheduleTaskInTxIfAbsent[T any, R any]( + tx server.PgTx, + taskFunction TaskFunction[T, R], + args T, + clientSession *session.ClientSession, + runOnce *RunOnceOption, + opts ...any, +) (scheduled bool, taskId server.Id) { + if runOnce == nil { + panic("ScheduleTaskInTxIfAbsent requires a non-nil runOnce key") + } + p := prepareTask(taskFunction, args, clientSession, append(opts, runOnce)...) + + claimTime := time.Time{} + + tag := server.RaisePgResult(tx.Exec( + clientSession.Ctx, + ` + INSERT INTO pending_task ( + task_id, + function_name, + args_json, + client_address, + client_by_jwt_json, + run_at, + run_once_key, + run_priority, + run_max_time_seconds, + claim_time, + release_time + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10) + ON CONFLICT (run_once_key) DO NOTHING + `, + p.taskId, + p.functionName, + p.argsJson, + clientSession.ClientAddress, + p.byJwtJson, + p.runAt, + p.runOnceKey, + p.priority, + p.maxTimeSeconds, + claimTime, + )) + scheduled = 0 < tag.RowsAffected() + return scheduled, p.taskId +} + +func ScheduleTaskIfAbsent[T any, R any]( + taskFunction TaskFunction[T, R], + args T, + clientSession *session.ClientSession, + runOnce *RunOnceOption, + opts ...any, +) (scheduled bool, taskId server.Id) { + server.Tx(clientSession.Ctx, func(tx server.PgTx) { + scheduled, taskId = ScheduleTaskInTxIfAbsent[T, R](tx, taskFunction, args, clientSession, runOnce, opts...) + }) return } diff --git a/task/task_test.go b/task/task_test.go index 562fa8aa..61289b4e 100644 --- a/task/task_test.go +++ b/task/task_test.go @@ -2,6 +2,7 @@ package task import ( "context" + "encoding/json" "errors" "fmt" mathrand "math/rand" @@ -205,3 +206,76 @@ func TestTask(t *testing.T) { assert.Equal(t, 0, len(ListFinishedTasks(ctx))) }) } + +type Work2Args struct { + Tag string +} + +type Work2Result struct{} + +func Work2( + work2 *Work2Args, + clientSession *session.ClientSession, +) (*Work2Result, error) { + return &Work2Result{}, nil +} + +// ScheduleTaskIfAbsent must atomically insert-or-detect-conflict on the +// run_once key: a first call inserts and reports scheduled == true; a second +// call with the SAME key while the first is still pending (unclaimed) must +// report scheduled == false and must NOT touch the first call's persisted +// args -- unlike plain ScheduleTask+RunOnce, whose ON CONFLICT DO UPDATE +// silently merges only timing/priority into the existing row while leaving +// (and thus never surfacing) the second call's args. +func TestScheduleTaskIfAbsent(t *testing.T) { + server.DefaultTestEnv().Run(t, func(t testing.TB) { + ctx := context.Background() + clientSession := session.Testing_CreateClientSession(ctx, nil) + defer clientSession.Cancel() + + key := RunOnce("test_schedule_task_if_absent", server.NewId()) + + scheduled, firstTaskId := ScheduleTaskIfAbsent( + Work2, + &Work2Args{Tag: "first"}, + clientSession, + key, + ) + assert.Equal(t, scheduled, true) + + // a second call with the same key, while the first is still pending, + // must be rejected -- not merged + scheduledAgain, _ := ScheduleTaskIfAbsent( + Work2, + &Work2Args{Tag: "second"}, + clientSession, + key, + ) + assert.Equal(t, scheduledAgain, false) + + // the persisted task must still be the FIRST call's args; the + // second call's args must never have been written anywhere + tasks := GetTasks(ctx, firstTaskId) + task, ok := tasks[firstTaskId] + if !ok { + t.Fatal("first task not found") + } + var args Work2Args + if err := json.Unmarshal([]byte(task.ArgsJson), &args); err != nil { + t.Fatal(err) + } + assert.Equal(t, args.Tag, "first") + + // once the pending row is gone (simulating the run finishing), the + // same key must be schedulable again + RemovePendingTask(ctx, firstTaskId) + + scheduledAfterClear, _ := ScheduleTaskIfAbsent( + Work2, + &Work2Args{Tag: "third"}, + clientSession, + key, + ) + assert.Equal(t, scheduledAfterClear, true) + }) +} diff --git a/taskworker/taskworker.go b/taskworker/taskworker.go index 8125be57..06708f47 100644 --- a/taskworker/taskworker.go +++ b/taskworker/taskworker.go @@ -6,6 +6,7 @@ import ( "github.com/urnetwork/server" "github.com/urnetwork/server/controller" + "github.com/urnetwork/server/model" "github.com/urnetwork/server/session" "github.com/urnetwork/server/task" "github.com/urnetwork/server/taskworker/work" @@ -151,6 +152,10 @@ func InitTaskWorker(ctx context.Context) *task.TaskWorker { work.SweepOrphanNetworkClientData, work.SweepOrphanNetworkClientDataPost, ), + task.NewTaskTargetWithPost( + model.RemoveNetworkClientsTask, + model.RemoveNetworkClientsTaskPost, + ), task.NewTaskTargetWithPost( work.SweepOrphanContractData, work.SweepOrphanContractDataPost,