From 02e2768abbb8e7d028dc7f5b2ac5fe909368f2e1 Mon Sep 17 00:00:00 2001 From: michel-laterman Date: Fri, 9 Jan 2026 10:10:18 -0600 Subject: [PATCH 1/3] Use req.Body readers to reduce allocs --- server/serverimpl.go | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/server/serverimpl.go b/server/serverimpl.go index 12048199..7c7e7dfc 100644 --- a/server/serverimpl.go +++ b/server/serverimpl.go @@ -245,11 +245,11 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co connectionCallbacks.OnConnected(reqCtx, agentConn) - sentCustomCapabilities := false + var sentCustomCapabilities sync.Once // Loop until fail to read from the WebSocket connection. for { - msgContext := context.Background() + msgContext := context.Background() // reqContext is cancelled when the ServerHTTP method is returned, for WebSockets connections this happens before this loop even starts. request := protobufs.AgentToServer{} // Block until the next message can be read. @@ -294,12 +294,11 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co if len(response.InstanceUid) == 0 { response.InstanceUid = request.InstanceUid } - if !sentCustomCapabilities { + sentCustomCapabilities.Do(func() { response.CustomCapabilities = &protobufs.CustomCapabilities{ Capabilities: s.settings.CustomCapabilities, } - sentCustomCapabilities = true - } + }) err = agentConn.Send(msgContext, response) if err != nil { @@ -311,8 +310,8 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co } } -func decompressGzip(data []byte) ([]byte, error) { - r, err := gzip.NewReader(bytes.NewBuffer(data)) +func decompressGzip(data io.Reader) ([]byte, error) { + r, err := gzip.NewReader(data) if err != nil { return nil, err } @@ -321,15 +320,16 @@ func decompressGzip(data []byte) ([]byte, error) { } func (s *server) readReqBody(req *http.Request) ([]byte, error) { - data, err := io.ReadAll(req.Body) - if err != nil { - return nil, err - } if req.Header.Get(headerContentEncoding) == contentEncodingGzip { - data, err = decompressGzip(data) + data, err = decompressGzip(req.Body) if err != nil { return nil, err } + return data, nil + } + data, err := io.ReadAll(req.Body) + if err != nil { + return nil, err } return data, nil } @@ -391,6 +391,7 @@ func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter } // Return the CustomCapabilities + // Note that unlike a WebSocket response, this is included in all HTTP responses. response.CustomCapabilities = &protobufs.CustomCapabilities{ Capabilities: s.settings.CustomCapabilities, } From 80bb0b057eb3c302684b62fadb579104e0d4f432 Mon Sep 17 00:00:00 2001 From: michel-laterman Date: Fri, 9 Jan 2026 10:27:33 -0600 Subject: [PATCH 2/3] Fix tests --- server/serverimpl.go | 2 +- server/serverimpl_test.go | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/server/serverimpl.go b/server/serverimpl.go index 7c7e7dfc..b0c34c29 100644 --- a/server/serverimpl.go +++ b/server/serverimpl.go @@ -321,7 +321,7 @@ func decompressGzip(data io.Reader) ([]byte, error) { func (s *server) readReqBody(req *http.Request) ([]byte, error) { if req.Header.Get(headerContentEncoding) == contentEncodingGzip { - data, err = decompressGzip(req.Body) + data, err := decompressGzip(req.Body) if err != nil { return nil, err } diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go index 74b1f1cc..62f48e75 100644 --- a/server/serverimpl_test.go +++ b/server/serverimpl_test.go @@ -964,11 +964,8 @@ func TestServerHonoursAcceptEncoding(t *testing.T) { // Verify the received message is what was sent. assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg)) - // Read Server's response. - b, err = io.ReadAll(resp.Body) - require.NoError(t, err) // Decompress the gzip response - b, err = decompressGzip(b) + b, err = decompressGzip(resp.Body) require.NoError(t, err) assert.EqualValues(t, http.StatusOK, resp.StatusCode) From 67b4b8e09d5f0fca3565d59c1da3a2849f2ac2c0 Mon Sep 17 00:00:00 2001 From: michel-laterman Date: Mon, 12 Jan 2026 08:33:05 -0800 Subject: [PATCH 3/3] Remove use of sync.Once --- server/serverimpl.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/server/serverimpl.go b/server/serverimpl.go index b0c34c29..f303df14 100644 --- a/server/serverimpl.go +++ b/server/serverimpl.go @@ -245,7 +245,7 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co connectionCallbacks.OnConnected(reqCtx, agentConn) - var sentCustomCapabilities sync.Once + sentCustomCapabilities := false // Loop until fail to read from the WebSocket connection. for { @@ -294,11 +294,12 @@ func (s *server) handleWSConnection(reqCtx context.Context, wsConn *websocket.Co if len(response.InstanceUid) == 0 { response.InstanceUid = request.InstanceUid } - sentCustomCapabilities.Do(func() { + if !sentCustomCapabilities { response.CustomCapabilities = &protobufs.CustomCapabilities{ Capabilities: s.settings.CustomCapabilities, } - }) + sentCustomCapabilities = true + } err = agentConn.Send(msgContext, response) if err != nil {