ssh/agent: support parallel signing using request pipelining Make NewClient automatically pipeline concurrent requests over its connection when the supplied io.ReadWriter also implements io.Closer. In this mode the client writes requests to the wire as soon as the write path is available and dispatches responses back to callers in FIFO order via a dedicated reader goroutine, instead of fully serializing each call. Up to 32 requests may be in flight on a single connection. This lets an agent that load-balances signing across multiple backend devices keep several of them busy concurrently without forcing callers to maintain their own connection pool. Aggregate throughput scales roughly linearly with the number of backend devices up to the in-flight cap; the protocol's in-order response requirement means slow requests still delay subsequent ones on the same connection, which is unchanged. The pipelined path requires io.Closer because, on a Write error, the background reader goroutine must be unblocked by closing the underlying connection; otherwise it would remain parked forever waiting for a response that will never arrive, leaking the goroutine and desynchronising the FIFO routing of responses for subsequent successful writes. When the supplied transport does not implement io.Closer, NewClient falls back to the previous fully-serialized behavior: a single in-flight call at a time, with no background goroutine. Fixes golang/go#78473 Change-Id: Icf14e5e8ca897506d68fb32c14fc72c774cd97b2 Reviewed-on: https://go-review.googlesource.com/c/crypto/+/768483 Reviewed-by: Cherry Mui <cherryyz@google.com> Reviewed-by: Dmitri Shuralyov <dmitshur@google.com> Reviewed-by: Filippo Valsorda <filippo@golang.org> LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/ssh/agent/client.go b/ssh/agent/client.go index 2fc2aa9..4276002 100644 --- a/ssh/agent/client.go +++ b/ssh/agent/client.go
@@ -26,6 +26,7 @@ "io" "math/big" "sync" + "sync/atomic" "golang.org/x/crypto/ssh" ) @@ -307,17 +308,50 @@ }, record.Rest, nil } +// pipelineMaxInFlight is the maximum number of outstanding requests the +// client will pipeline to the agent before applying backpressure. +const pipelineMaxInFlight = 32 + // client is a client for an ssh-agent process. +// +// Exactly one of pipeline / (mu, conn) is set, chosen by NewClient +// based on whether the underlying transport implements io.Closer. type client struct { - // conn is typically a *net.UnixConn + // pipeline, if non-nil, dispatches requests over a pipelined + // connection: requests are written as soon as the wire is + // available and responses are routed back to per-call reply + // channels in FIFO order by a background reader goroutine. + pipeline *pipeline + + // mu and conn are used in fully-serialized mode, when the + // transport does not implement io.Closer. Each call takes mu, + // writes its request, reads the matching response, and releases + // mu before returning. There is no background goroutine. + mu sync.Mutex conn io.ReadWriter - // mu is used to prevent concurrent access to the agent - mu sync.Mutex } // NewClient returns an Agent that talks to an ssh-agent process over // the given connection. +// +// If rw also implements io.Closer (like *net.UnixConn and ssh.Channel +// do), the returned client pipelines concurrent requests over the +// connection: callers can issue Sign and other operations from +// multiple goroutines and they will be written to the agent as soon +// as the wire is available, rather than waiting for the previous +// responses. The ssh-agent protocol still requires responses to be +// returned in request order, so a slow request delays subsequent +// responses on the same connection (head-of-line blocking). +// +// Pipelining requires io.Closer because, on a Write error, the +// background reader goroutine must be unblocked by closing the +// underlying connection. When rw does not implement io.Closer +// this is not possible, so NewClient falls back to fully +// serializing each request: a single in-flight call at a time. func NewClient(rw io.ReadWriter) ExtendedAgent { + if rwc, ok := rw.(io.ReadWriteCloser); ok { + return &client{pipeline: newPipeline(rwc)} + } return &client{conn: rw} } @@ -340,6 +374,16 @@ // bytes of the response are returned; no unmarshalling is // performed on the response. func (c *client) callRaw(req []byte) (reply []byte, err error) { + if c.pipeline != nil { + return c.pipeline.call(req) + } + return c.serialCall(req) +} + +// serialCall implements the fully-serialized request/response path +// used when the transport is not an io.Closer. It writes req under mu +// and reads the matching response before returning. +func (c *client) serialCall(req []byte) (reply []byte, err error) { c.mu.Lock() defer c.mu.Unlock() @@ -861,3 +905,170 @@ return buf, nil } + +// pipelineResult carries either a raw agent reply or an error back to a +// caller waiting on the response channel. +type pipelineResult struct { + reply []byte + err error +} + +// pipeline implements request pipelining over a single agent connection. +// +// Writers serialize on writeMu to both register a reply channel in the +// pending FIFO queue and write the request bytes on the wire; the two +// must be atomic so the queue order matches the wire order. A single +// reader goroutine decodes responses from the connection and dispatches +// each one to the channel at the head of the queue. +// +// pending is a chan-of-chan acting as a FIFO queue with a fixed +// capacity of pipelineMaxInFlight. The outer channel provides ordering +// (reads happen in send order) and natural backpressure (a full queue +// blocks new writers). Each inner channel is buffered with capacity +// one and is sent to exactly once: either by the reader goroutine +// with the agent reply, or by shutdown with the terminal error during +// drain. The cap-one buffer makes the producer's send non-blocking, +// so the reader and shutdown never have to wait for the caller to be +// scheduled on the receive. +// +// When the reader goroutine exits (on read error or protocol +// violation), it closes exitCh to wake any writer blocked on the +// pending queue, then serializes with any in-flight writer to close +// the pending channel, and finally drains the remaining entries +// delivering the terminal error to each waiting caller. The +// pipeline relies on conn implementing io.Closer so a writer that +// hits a Write error can close the connection to unblock the reader +// goroutine; NewClient is responsible for only constructing a +// pipeline when this guarantee holds. +type pipeline struct { + conn io.ReadWriteCloser + + writeMu sync.Mutex + // pending is the FIFO queue of reply channels with capacity + // pipelineMaxInFlight. See type-level documentation. + pending chan chan pipelineResult + exitCh chan struct{} + + // err carries the terminal error to callers blocked on a closed + // pipeline. It is stored exactly once by the reader goroutine + // before exitCh is closed; every read happens after observing + // exitCh closed, so the load synchronises through the close and + // is guaranteed to return the stored value (never nil). + err atomic.Pointer[error] +} + +func newPipeline(conn io.ReadWriteCloser) *pipeline { + p := &pipeline{ + conn: conn, + pending: make(chan chan pipelineResult, pipelineMaxInFlight), + exitCh: make(chan struct{}), + } + go p.readLoop() + return p +} + +// readLoop decodes responses from conn and dispatches them in FIFO order +// to reply channels in pending. On any failure it invokes shutdown. +func (p *pipeline) readLoop() { + var finalErr error + for { + var sizeBuf [4]byte + if _, err := io.ReadFull(p.conn, sizeBuf[:]); err != nil { + finalErr = err + break + } + respSize := binary.BigEndian.Uint32(sizeBuf[:]) + if respSize > maxAgentResponseBytes { + finalErr = errors.New("response too large") + break + } + buf := make([]byte, respSize) + if _, err := io.ReadFull(p.conn, buf); err != nil { + finalErr = err + break + } + // Successful writes always enqueue before sending bytes, so + // pending has a waiting channel for this response. + ch := <-p.pending + // The reply channel is buffered with capacity 1 and is only + // ever written to once, so this send cannot block. + ch <- pipelineResult{reply: buf} + } + p.shutdown(clientErr(finalErr)) +} + +// shutdown is called exactly once, from readLoop, when the reader is +// terminating. It unblocks pending writers and fails all in-flight +// requests with finalErr. +func (p *pipeline) shutdown(finalErr error) { + // Publish the terminal error before closing exitCh so any + // writer that subsequently observes exitCh closed sees err. + p.err.Store(&finalErr) + + // Wake any writer blocked waiting for a slot in the pending queue. + close(p.exitCh) + + // Wait for any writer currently inside its critical section to + // complete. After this lock, no new writer can reach the send on + // pending: they will observe exitCh closed in the select and bail + // out before attempting the send. + p.writeMu.Lock() + close(p.pending) + p.writeMu.Unlock() + + // Drain entries that were enqueued but never answered, delivering + // the terminal error to their waiting callers. The reply channels + // are buffered (cap 1) and written to exactly once, so these sends + // cannot block. + for ch := range p.pending { + ch <- pipelineResult{err: finalErr} + } +} + +// call sends req to the agent and returns the matching raw response. +func (p *pipeline) call(req []byte) ([]byte, error) { + replyCh := make(chan pipelineResult, 1) + + p.writeMu.Lock() + + // Priority check: if the reader has already finished shutdown, + // pending is closed and sending to it would panic. Bail out now. + // Once we pass this check while holding writeMu, shutdown cannot + // complete close(pending) until we release writeMu, so the send + // below is safe against concurrent closure. + select { + case <-p.exitCh: + p.writeMu.Unlock() + return nil, *p.err.Load() + default: + } + + // Enqueue the reply channel before writing the request, so FIFO + // order on the wire matches FIFO order in the pending queue. The + // exitCh arm handles the case where the reader errors while we + // block on a full queue. + select { + case p.pending <- replyCh: + case <-p.exitCh: + p.writeMu.Unlock() + return nil, *p.err.Load() + } + + msg := make([]byte, 4+len(req)) + binary.BigEndian.PutUint32(msg, uint32(len(req))) + copy(msg[4:], req) + _, werr := p.conn.Write(msg) + p.writeMu.Unlock() + + if werr != nil { + // The connection is in an undefined state. Close it so the + // reader unblocks promptly and triggers shutdown for every + // other in-flight caller. NewClient guarantees conn is a + // real io.Closer when the pipeline is in use. + p.conn.Close() + return nil, clientErr(werr) + } + + res := <-replyCh + return res.reply, res.err +}
diff --git a/ssh/agent/client_pipeline_test.go b/ssh/agent/client_pipeline_test.go new file mode 100644 index 0000000..d453f2c --- /dev/null +++ b/ssh/agent/client_pipeline_test.go
@@ -0,0 +1,649 @@ +// Copyright 2026 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package agent + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestPipelineParallelSigns(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c1.Close() + defer c2.Close() + + keyring := NewKeyring() + if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil { + t.Fatal(err) + } + go ServeAgent(keyring, c2) + + client := NewClient(c1) + pubKey := testPublicKeys["rsa"] + + const N = 64 + var wg sync.WaitGroup + errCh := make(chan error, N) + for i := range N { + wg.Go(func() { + data := []byte{byte(i), byte(i >> 8)} + sig, err := client.Sign(pubKey, data) + if err != nil { + errCh <- err + return + } + if err := pubKey.Verify(data, sig); err != nil { + errCh <- err + } + }) + } + wg.Wait() + close(errCh) + for err := range errCh { + t.Error(err) + } +} + +func TestPipelineBackpressure(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + + const total = pipelineMaxInFlight + 8 + + // Fake agent: the reader drains requests from the wire independently + // of response production, so the write path is never self-blocked. + // The responder only writes a canned success reply once unblocked + // via gate. + var reqsRead int64 + gate := make(chan struct{}, total) + pending := make(chan struct{}, total+pipelineMaxInFlight) + readerDone := make(chan struct{}) + responderDone := make(chan struct{}) + go func() { + defer close(readerDone) + defer close(pending) + for { + var sizeBuf [4]byte + if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil { + return + } + size := binary.BigEndian.Uint32(sizeBuf[:]) + req := make([]byte, size) + if _, err := io.ReadFull(c2, req); err != nil { + return + } + atomic.AddInt64(&reqsRead, 1) + pending <- struct{}{} + } + }() + go func() { + defer close(responderDone) + for range pending { + <-gate + var out [5]byte + binary.BigEndian.PutUint32(out[:4], 1) + out[4] = agentSuccess + if _, err := c2.Write(out[:]); err != nil { + return + } + } + }() + defer func() { + c1.Close() + c2.Close() + <-readerDone + <-responderDone + }() + + client := NewClient(c1) + + var wg sync.WaitGroup + errs := make(chan error, total) + for range total { + wg.Go(func() { + // RemoveAll maps to a simpleCall expecting agentSuccess. + if err := client.RemoveAll(); err != nil { + errs <- err + } + }) + } + + // Wait for the client to fill the pipeline. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight { + break + } + time.Sleep(5 * time.Millisecond) + } + if got := atomic.LoadInt64(&reqsRead); got != pipelineMaxInFlight { + t.Errorf("requests in flight before any response: got %d, want %d", got, pipelineMaxInFlight) + } + + // Release all responses; every caller should eventually complete. + for range total { + gate <- struct{}{} + } + wg.Wait() + close(errs) + for err := range errs { + t.Error(err) + } + + if got := atomic.LoadInt64(&reqsRead); got != total { + t.Errorf("total requests received: got %d, want %d", got, total) + } +} + +func TestPipelineConnCloseFailsInFlight(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + + // Server that reads requests but never replies, so all calls are + // parked waiting for responses. + serverExit := make(chan struct{}) + go func() { + defer close(serverExit) + buf := make([]byte, 4096) + for { + if _, err := c2.Read(buf); err != nil { + return + } + } + }() + + client := NewClient(c1) + + const N = 8 + var wg sync.WaitGroup + errs := make(chan error, N) + for range N { + wg.Go(func() { + errs <- client.RemoveAll() + }) + } + + // Give callers time to enter the pipeline, then tear down the conn. + time.Sleep(50 * time.Millisecond) + c1.Close() + c2.Close() + + // All callers must return promptly with an error. + doneCh := make(chan struct{}) + go func() { + wg.Wait() + close(doneCh) + }() + select { + case <-doneCh: + case <-time.After(5 * time.Second): + t.Fatal("pipelined callers did not complete after conn close") + } + + close(errs) + for err := range errs { + if err == nil { + t.Error("expected error after conn close, got nil") + } + } + + // Subsequent calls also fail fast rather than hanging. + errCh := make(chan error, 1) + go func() { errCh <- client.RemoveAll() }() + select { + case err := <-errCh: + if err == nil { + t.Error("expected error on post-close call, got nil") + } + case <-time.After(2 * time.Second): + t.Fatal("post-close call hung") + } + + <-serverExit +} + +// readWriter exposes only io.Reader and io.Writer, hiding any Close +// method on the wrapped value so NewClient cannot take the pipelined +// path. +type readWriter struct { + io.Reader + io.Writer +} + +func TestNewClientSerialFallback(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c1.Close() + defer c2.Close() + + keyring := NewKeyring() + if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil { + t.Fatal(err) + } + go ServeAgent(keyring, c2) + + rw := readWriter{Reader: c1, Writer: c1} + ag := NewClient(rw) + + ci, ok := ag.(*client) + if !ok { + t.Fatalf("NewClient: got %T, want *client", ag) + } + if ci.pipeline != nil { + t.Error("NewClient with non-Closer rw: pipeline must be nil") + } + if ci.conn == nil { + t.Error("NewClient with non-Closer rw: conn must be set") + } + + // The serialized path correctly handles a request/response cycle. + keys, err := ag.List() + if err != nil { + t.Fatalf("List: %v", err) + } + if len(keys) != 1 { + t.Errorf("List: got %d keys, want 1", len(keys)) + } +} + +func TestSerialFallbackParallelSigns(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c1.Close() + defer c2.Close() + + keyring := NewKeyring() + if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil { + t.Fatal(err) + } + go ServeAgent(keyring, c2) + + rw := readWriter{Reader: c1, Writer: c1} + ag := NewClient(rw) + + if ci, ok := ag.(*client); !ok { + t.Fatalf("NewClient: got %T, want *client", ag) + } else if ci.pipeline != nil { + t.Fatal("pipeline must be nil for non-Closer rw") + } + + pubKey := testPublicKeys["rsa"] + + const N = 64 + var wg sync.WaitGroup + errCh := make(chan error, N) + for i := range N { + wg.Go(func() { + data := []byte{byte(i), byte(i >> 8)} + sig, err := ag.Sign(pubKey, data) + if err != nil { + errCh <- err + return + } + if err := pubKey.Verify(data, sig); err != nil { + errCh <- err + } + }) + } + wg.Wait() + close(errCh) + for err := range errCh { + t.Error(err) + } +} + +// writeFailConn wraps a net.Conn and forces Write to return an error when +// the fail flag is set. Close delegates to the underlying conn so the test +// can observe whether the pipeline closed it. +type writeFailConn struct { + net.Conn + fail atomic.Bool +} + +func (c *writeFailConn) Write(p []byte) (int, error) { + if c.fail.Load() { + return 0, errors.New("synthetic write failure") + } + return c.Conn.Write(p) +} + +func TestPipelineWriteErrorClosesConn(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c2.Close() + + // Peer that reads and discards, exiting when the conn closes. + serverExit := make(chan struct{}) + go func() { + defer close(serverExit) + io.Copy(io.Discard, c2) + }() + + wrapper := &writeFailConn{Conn: c1} + wrapper.fail.Store(true) + + client := NewClient(wrapper) + + errCh := make(chan error, 1) + go func() { errCh <- client.RemoveAll() }() + + select { + case err := <-errCh: + if err == nil { + t.Fatal("expected error on forced write failure, got nil") + } + // The caller whose Write failed must receive the actual + // write error, not the reader's "closed connection" error + // from the post-shutdown drain. + if !strings.Contains(err.Error(), "synthetic write failure") { + t.Errorf("got %q, want error containing %q", err, "synthetic write failure") + } + case <-time.After(2 * time.Second): + t.Fatal("RemoveAll hung after write failure") + } + + // The pipeline must have closed the underlying conn, so the peer + // goroutine returns from its Read. + select { + case <-serverExit: + case <-time.After(2 * time.Second): + t.Fatal("connection was not closed after write failure") + } + + // After shutdown, subsequent calls must also fail promptly. + errCh2 := make(chan error, 1) + go func() { errCh2 <- client.RemoveAll() }() + select { + case err := <-errCh2: + if err == nil { + t.Error("expected error on post-shutdown call, got nil") + } + case <-time.After(2 * time.Second): + t.Fatal("post-shutdown call hung") + } +} + +func TestPipelineResponseTooLarge(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c1.Close() + defer c2.Close() + + // Server: read the request, then reply with a size header that + // exceeds maxAgentResponseBytes. The size check fires before the + // body is read, so no body needs to follow. + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + var sizeBuf [4]byte + if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil { + return + } + size := binary.BigEndian.Uint32(sizeBuf[:]) + if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil { + return + } + var out [4]byte + binary.BigEndian.PutUint32(out[:], maxAgentResponseBytes+1) + c2.Write(out[:]) + }() + + client := NewClient(c1) + + errCh := make(chan error, 1) + go func() { errCh <- client.RemoveAll() }() + select { + case err := <-errCh: + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "response too large") { + t.Errorf("got %q, want error containing %q", err, "response too large") + } + case <-time.After(2 * time.Second): + t.Fatal("call hung") + } + + // Subsequent calls must also fail: the pipeline is shut down. + errCh2 := make(chan error, 1) + go func() { errCh2 <- client.RemoveAll() }() + select { + case err := <-errCh2: + if err == nil { + t.Error("expected error on post-shutdown call, got nil") + } + case <-time.After(2 * time.Second): + t.Fatal("post-shutdown call hung") + } + + <-serverDone +} + +func TestPipelineShutdownWithBlockedWriter(t *testing.T) { + c1, c2, err := netPipe() + if err != nil { + t.Fatalf("netPipe: %v", err) + } + defer c1.Close() + + // Server drains requests off the wire so writes never block on + // backpressure from the transport, but never replies. The reader + // goroutine therefore stays parked on Read, never dequeues from + // pending, and the queue fills up. + var reqsRead int64 + serverExit := make(chan struct{}) + go func() { + defer close(serverExit) + for { + var sizeBuf [4]byte + if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil { + return + } + size := binary.BigEndian.Uint32(sizeBuf[:]) + if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil { + return + } + atomic.AddInt64(&reqsRead, 1) + } + }() + + client := NewClient(c1) + + const total = pipelineMaxInFlight + 4 + var wg sync.WaitGroup + errs := make(chan error, total) + for range total { + wg.Go(func() { + errs <- client.RemoveAll() + }) + } + + // Wait for the queue to fill: once pipelineMaxInFlight requests + // are on the wire, pending is at capacity and the +1 writer is + // blocked in the second-select arm; the rest are queued behind + // writeMu. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight { + break + } + time.Sleep(5 * time.Millisecond) + } + if got := atomic.LoadInt64(&reqsRead); got < pipelineMaxInFlight { + t.Fatalf("requests on wire: got %d, want >= %d", got, pipelineMaxInFlight) + } + // Give the +1 writer a moment to actually park in the second + // select before we trigger shutdown. + time.Sleep(50 * time.Millisecond) + + // Trigger shutdown: closing the server side makes the client's + // reader hit EOF, which calls shutdown(...). + c2.Close() + + doneCh := make(chan struct{}) + go func() { + wg.Wait() + close(doneCh) + }() + select { + case <-doneCh: + case <-time.After(5 * time.Second): + t.Fatal("blocked callers did not complete after shutdown") + } + + close(errs) + for err := range errs { + if err == nil { + t.Error("expected error after shutdown, got nil") + } + } + + <-serverExit +} + +// multiDeviceAgent simulates an ssh-agent backed by numDevices signing +// devices, each of which takes latency to process a request. Incoming +// requests are round-robin dispatched to devices, and responses are +// serialized back to the wire in request order as the protocol requires. +// +// The simulation does not emit a semantically meaningful reply; it just +// writes a 1-byte SSH_AGENT_SUCCESS after the simulated latency has +// elapsed. This is enough to exercise callRaw/simpleCall (RemoveAll). +func multiDeviceAgent(t testing.TB, conn io.ReadWriteCloser, numDevices int, latency time.Duration) (done <-chan struct{}) { + t.Helper() + + type job struct { + ready chan struct{} + } + + devices := make([]chan job, numDevices) + var devicesWG sync.WaitGroup + for i := range numDevices { + devices[i] = make(chan job, 1024) + ch := devices[i] + devicesWG.Go(func() { + for j := range ch { + time.Sleep(latency) + close(j.ready) + } + }) + } + + // Order-preserving responder: responses are written in the same + // order the requests were read, regardless of which device finished + // first. + orderCh := make(chan chan struct{}, 1024) + responderDone := make(chan struct{}) + go func() { + defer close(responderDone) + for ready := range orderCh { + <-ready + var out [5]byte + binary.BigEndian.PutUint32(out[:4], 1) + out[4] = agentSuccess + if _, err := conn.Write(out[:]); err != nil { + return + } + } + }() + + finished := make(chan struct{}) + go func() { + defer close(finished) + defer func() { + close(orderCh) + <-responderDone + for _, d := range devices { + close(d) + } + devicesWG.Wait() + }() + next := 0 + for { + var sizeBuf [4]byte + if _, err := io.ReadFull(conn, sizeBuf[:]); err != nil { + return + } + size := binary.BigEndian.Uint32(sizeBuf[:]) + if _, err := io.CopyN(io.Discard, conn, int64(size)); err != nil { + return + } + ready := make(chan struct{}) + devices[next%numDevices] <- job{ready: ready} + next++ + orderCh <- ready + } + }() + + return finished +} + +// BenchmarkPipelineMultiDevice measures throughput of the pipelined +// agent client against a simulated multi-device agent, varying the +// number of backend devices. +// +// Setup: each simulated signing device takes 2ms per request. The +// protocol serializes responses in request order, but pipelining lets +// the client keep multiple devices busy concurrently; throughput +// should scale roughly linearly with numDevices up to the pipeline +// depth (pipelineMaxInFlight). +// +// Run with: go test -benchmem -bench BenchmarkPipelineMultiDevice -benchtime=2s ./ssh/agent/ +func BenchmarkPipelineMultiDevice(b *testing.B) { + const ( + deviceLatency = 2 * time.Millisecond + parallelism = 32 + ) + + for _, numDevices := range []int{1, 2, 4, 8, 16} { + b.Run(fmt.Sprintf("Devices=%d", numDevices), func(b *testing.B) { + c1, c2, err := netPipe() + if err != nil { + b.Fatalf("netPipe: %v", err) + } + serverDone := multiDeviceAgent(b, c2, numDevices, deviceLatency) + + client := NewClient(c1) + + b.ResetTimer() + b.SetParallelism(parallelism) + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + if err := client.RemoveAll(); err != nil { + b.Fatal(err) + } + } + }) + b.StopTimer() + + c1.Close() + c2.Close() + <-serverDone + }) + } +}