Pool marshal buffers for Duo Workflow action writes

Igor Drozdovrequested to merge
fix/pool-duo-workflow-marshal-buffers into master

What does this MR do and why?

This removes a per-connection buffer in the Duo Workflow WebSocket handler in Workhorse, and replaces it with a shared pool.

wsManager had a buf []byte field. It was used to marshal outbound protobuf Action messages into JSON before writing them to the client's WebSocket. The buffer grew to the largest action the connection had seen, and then stayed that large for the whole life of the connection. These connections stay open for over an hour.

Production actions are routinely 1-2MB. Peak concurrency in gprd is about 950 simultaneous connections. So the fleet pinned roughly 1.4MB per open connection, over a gigabyte in total, which brought the container close to its 2GB memory limit.

The per-connection field is removed. WriteAction now borrows a buffer from a package-level sync.Pool and returns it when the function returns. Retention now scales with the number of writes happening at the same time, instead of the number of connections that are open. That matters because connections are idle nearly all the time, waiting on model and tool latency measured in seconds, while marshaling and writing takes milliseconds. The pool keeps the grown capacity, so repeated large actions do not reallocate. And because sync.Pool is drained by the garbage collector, memory is returned under pressure instead of being held at a high-water mark.

Measurements

Local benchmarks on an Apple M3 Max, 500 connections each writing one 1.42MB action.

Measurement Before After
Retained per connection 1,458,176 bytes 24 bytes
Steady-state write cost 2,207 B/op 1,469 B/op, 19 allocs/op

A linear regression of production heap against the gitlab_workhorse_duo_workflow_connections_open gauge gives a marginal cost of about 1.63MB per connection. That is close to the 1.39MB measured locally, which confirms the buffer was the dominant per-connection cost.

Tradeoff

The buffer starts at 4KB. In a shared pool, every pool miss has to grow from 4KB again, which is roughly 22 reallocations to reach a 1.4MB action. Misses cluster right after garbage collection, because the collector drains the pool. So this trades a steady, always-present live heap for some allocation churn that is correlated with GC. It is still a clear net win against pinning 950 times 1.4MB, but tuning the initial size against the production size distribution is a sensible follow-up.

Test change

mockWebSocketConn.WriteMessage now copies the slice it is given. The real gorilla/websocket WriteMessage writes synchronously and keeps no reference, but the mock retained it, which is unsafe once the buffer is recycled.

Verification

Package tests pass, including 35 runs with -shuffle=on and 4 runs with -race (zero data races). gofmt and go vet are clean.

One thing so reviewers are not surprised: TestHeaderParsing in this package is an order-dependent pre-existing failure. It fails even when run alone, because it depends on another test seeding a global JWT secret path. This MR does not touch it.

References

Screenshots or screen recordings

Not applicable. This is a backend memory change with no user interface.

How to set up and validate locally

Run all commands from the workhorse/ directory.

  1. Start GDK Redis first: gdk start redis. Without it, unrelated tests in this package fail on a missing socket.

  2. Run the package tests:

    go test ./internal/ai_assist/duoworkflow/ -count=1

  3. Measure retained memory per connection and look at the B-retained/conn metric:

    go test ./internal/ai_assist/duoworkflow/ -run XXX -bench 'BenchmarkConnectionsRetainedBytes' -benchtime=1x

  4. Confirm the per-write cost did not regress:

    go test ./internal/ai_assist/duoworkflow/ -run XXX -bench 'BenchmarkWriteActionLarge' -benchtime=2000x -benchmem

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Merge request reports