ba-proxy-agent: close idle connections to mitigate memory leaks

The ba-proxy-agent currently experiences increasing memory consumption,
leading to daily or weekly restarts. Previous mitigation attempts were
insufficient, and the issue has been isolated to high memory usage on
the agent connection side.

This CL mitigates the issue by enforcing a timeout on idle connections.
Since the root cause remains elusive, forcefully closing unused
connections prevents memory accumulation from persistent links.

Implementation details:
* Introduced `lastActivityTime` property to agent connections, which
    updates upon usage.
* Added a background routine to monitor connection activity.
* Configured the routine to explicitly close connections that remain
    idle for more than 30 seconds.

Author: hareeshv

agent/agent.go

@@ -73,6 +73,7 @@ var (
 	injectBanner              = flag.String("inject-banner", "", "HTML snippet to inject in served webpages")
 	bannerHeight              = flag.String("banner-height", "40px", "Height of the injected banner. This is ignored if no banner is set.")
 	shimWebsockets            = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
+	websocketShimTimeout      = flag.Duration("websocket-shim-timeout", 60*time.Minute, "Timeout for websocket shim connections to expire due to inactivity.")
 	shimPath                  = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
 	healthCheckPath           = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
 	healthCheckFreq           = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
@@ -126,7 +127,8 @@ func hostProxy(ctx context.Context, host, shimPath string, injectShimCode, force
 		// restricted to a path prefix not equal to "/" will fail for websocket open requests. Passing in the
 		// sessionHandler twice allows the websocket handler to ensure that cookies are applied based on the
 		// correct, restored path.
-		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, *enableWebsocketsInjection, sessionLRU.SessionHandler, metricHandler)
+		h, err = websockets.Proxy(ctx, h, host, shimPath, *rewriteWebsocketHost, *enableWebsocketsInjection, sessionLRU.SessionHandler,
+		                          metricHandler, *websocketShimTimeout)
 		if injectShimCode {
 			shimFunc, err := websockets.ShimBody(shimPath)
 			if err != nil {

agent/websockets/connection.go

@@ -17,16 +17,16 @@ limitations under the License.
 package websockets
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"log"
 	"net/http"
+	"sync"
 	"time"
 
-	"context"
-
 	"github.com/gorilla/websocket"
 )
 
@@ -57,12 +57,14 @@ func (m *message) Serialize(version int) interface{} {
 // and encapsulates it in an API that is a little more amenable to how the server side
 // of our websocket shim is implemented.
 type Connection struct {
-	done            func() <-chan struct{}
-	cancel          context.CancelFunc
-	clientMessages  chan *message
-	serverMessages  chan *message
-	protocolVersion int
-	subprotocol     string
+	done             func() <-chan struct{}
+	cancel           context.CancelFunc
+	clientMessages   chan *message
+	serverMessages   chan *message
+	protocolVersion  int
+	subprotocol      string
+	mu               sync.Mutex
+	lastActivityTime time.Time
 }
 
 // This map defines the set of headers that should be stripped from the WS request, as they
@@ -87,6 +89,20 @@ func stripWSHeader(header http.Header) http.Header {
 	return result
 }
 
+// updateActivity updates the last activity timestamp.
+func (conn *Connection) updateActivity() {
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+	conn.lastActivityTime = time.Now()
+}
+
+// lastActivity returns the last activity timestamp.
+func (conn *Connection) lastActivity() time.Time {
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+	return conn.lastActivityTime
+}
+
 // NewConnection creates and returns a new Connection.
 func NewConnection(ctx context.Context, targetURL string, header http.Header, errCallback func(err error)) (*Connection, error) {
 	ctx, cancel := context.WithCancel(ctx)
@@ -162,11 +178,12 @@ func NewConnection(ctx context.Context, targetURL string, header http.Header, er
 		}
 	}()
 	return &Connection{
-		done:           ctx.Done,
-		cancel:         cancel,
-		clientMessages: clientMessages,
-		serverMessages: serverMessages,
-		subprotocol: serverConn.Subprotocol(),
+		done:             ctx.Done,
+		cancel:           cancel,
+		clientMessages:   clientMessages,
+		serverMessages:   serverMessages,
+		subprotocol:      serverConn.Subprotocol(),
+		lastActivityTime: time.Now(),
 	}, nil
 }
 
@@ -184,6 +201,7 @@ func (conn *Connection) Close() {
 //
 // The returned error value is non-nill if the connection has been closed.
 func (conn *Connection) SendClientMessage(msg interface{}, injectionEnabled bool, injectedHeaders map[string]string) error {
+	conn.updateActivity()
 	var clientMessage *message
 	if textMsg, ok := msg.(string); ok {
 		clientMessage = &message{
@@ -236,7 +254,9 @@ func (conn *Connection) SendClientMessage(msg interface{}, injectionEnabled bool
 //
 // The returned []string value is nil if the error is non-nil, or if the method
 // times out while waiting for a server message.
-func (conn *Connection) ReadServerMessages() ([]interface{}, error) {
+func (conn *Connection) ReadServerMessages(readTimeout time.Duration) ([]interface{}, error) {
+	conn.updateActivity()
+	defer conn.updateActivity()
 	var msgs []interface{}
 	select {
 	case serverMsg, ok := <-conn.serverMessages:
@@ -257,7 +277,7 @@ func (conn *Connection) ReadServerMessages() ([]interface{}, error) {
 				return msgs, nil
 			}
 		}
-	case <-time.After(time.Second * 20):
+	case <-time.After(readTimeout):
 		return nil, nil
 	}
 }

agent/websockets/shim.go

@@ -18,6 +18,7 @@ package websockets
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -31,8 +32,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"text/template"
+	"time"
 
-	"context"
 	"github.com/google/inverting-proxy/agent/metrics"
 )
 
@@ -320,9 +321,33 @@ func (c *connectionErrorHandler) ReportError(err error) {
 	}
 }
 
-func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost bool, openWebsocketWrapper func(http.Handler, *metrics.MetricHandler) http.Handler, enableWebsocketInjection bool, metricHandler *metrics.MetricHandler) http.Handler {
+func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost bool, openWebsocketWrapper func(http.Handler, *metrics.MetricHandler) http.Handler, enableWebsocketInjection bool, metricHandler *metrics.MetricHandler, timeout time.Duration) http.Handler {
 	var connections sync.Map
 	var sessionCount uint64
+
+	// Background goroutine to clean up inactive websocket shim connections.
+	go func() {
+		ticker := time.NewTicker(min(timeout, 30*time.Second))
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				connections.Range(func(key, value any) bool {
+					sessionID := key.(string)
+					conn := value.(*Connection)
+					if time.Since(conn.lastActivity()) > timeout {
+						log.Printf("Closing inactive websocket shim session %q after timeout", sessionID)
+						conn.Close()
+						connections.Delete(sessionID)
+					}
+					return true // Continue iteration
+				})
+			}
+		}
+	}()
+
 	mux := http.NewServeMux()
 	errorHandler := &connectionErrorHandler{}
 	openWebsocketHandler := openWebsocketWrapper(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -351,9 +376,9 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 			}
 		}
 		resp := &sessionMessage{
-			ID:      sessionID,
-			Message: targetURL.String(),
-			Version: conn.protocolVersion,
+			ID:          sessionID,
+			Message:     targetURL.String(),
+			Version:     conn.protocolVersion,
 			Subprotocol: conn.Subprotocol(),
 		}
 		respBytes, err := json.Marshal(resp)
@@ -512,7 +537,7 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 			metricHandler.WriteResponseCodeMetric(statusCode)
 			return
 		}
-		serverMsgs, err := conn.ReadServerMessages()
+		serverMsgs, err := conn.ReadServerMessages(min(20*time.Second, timeout/2))
 		if err != nil {
 			statusCode := http.StatusBadRequest
 			errorMessage := fmt.Sprintf("attempt to read data from a closed session: %q", msg.ID)
@@ -548,11 +573,11 @@ func createShimChannel(ctx context.Context, host, shimPath string, rewriteHost b
 // openWebsocketWrapper is a http.Handler wrapper function that is invoked on websocket open requests after the original
 // targetURL of the request is restored. It must call the wrapped http.Handler with which it is created after it
 // is finished processing the request.
-func Proxy(ctx context.Context, wrapped http.Handler, host, shimPath string, rewriteHost, enableWebsocketInjection bool, openWebsocketWrapper func(wrapped http.Handler, metricHandler *metrics.MetricHandler) http.Handler, metricHandler *metrics.MetricHandler) (http.Handler, error) {
+func Proxy(ctx context.Context, wrapped http.Handler, host, shimPath string, rewriteHost, enableWebsocketInjection bool, openWebsocketWrapper func(wrapped http.Handler, metricHandler *metrics.MetricHandler) http.Handler, metricHandler *metrics.MetricHandler, timeout time.Duration) (http.Handler, error) {
 	mux := http.NewServeMux()
 	if shimPath != "" {
 		shimPath = path.Clean("/"+shimPath) + "/"
-		shimServer := createShimChannel(ctx, host, shimPath, rewriteHost, openWebsocketWrapper, enableWebsocketInjection, metricHandler)
+		shimServer := createShimChannel(ctx, host, shimPath, rewriteHost, openWebsocketWrapper, enableWebsocketInjection, metricHandler, timeout)
 		mux.Handle(shimPath, shimServer)
 	}
 	mux.Handle("/", wrapped)

agent/websockets/websockets_test.go

@@ -23,13 +23,15 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"path"
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
@@ -239,7 +241,7 @@ func TestShimHandlers(t *testing.T) {
 	openWrapper := func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler {
 		return h
 	}
-	p, err := Proxy(context.Background(), h, serverURL.Host, testShimPath, false, false, openWrapper, nil)
+	p, err := Proxy(context.Background(), h, serverURL.Host, testShimPath, false, false, openWrapper, nil, 60*time.Second)
 	if err != nil {
 		t.Fatalf("Failure creating the websocket shim proxy: %+v", err)
 	}
@@ -354,3 +356,107 @@ func TestShimHandlers(t *testing.T) {
 		}
 	}
 }
+
+func TestShimPolling(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	// Setup a fake backend that accepts websocket connections but sends no messages.
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		upgrader := websocket.Upgrader{}
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			t.Logf("Failed to upgrade websocket: %v", err)
+			return
+		}
+		defer conn.Close()
+		// Keep connection open until client closes it.
+		for {
+			if _, _, err := conn.NextReader(); err != nil {
+				break
+			}
+		}
+	}))
+	defer backend.Close()
+
+	backendURL, err := url.Parse(backend.URL)
+	if err != nil {
+		t.Fatalf("Failed to parse backend URL: %v", err)
+	}
+
+	shimPath := "/shim/"
+	idleTimeout := 3 * time.Second
+	shim := createShimChannel(
+		ctx,
+		backendURL.Host,
+		shimPath,
+		false,
+		func(h http.Handler, m *metrics.MetricHandler) http.Handler { return h },
+		false,
+		nil,
+		idleTimeout,
+	)
+	shimServer := httptest.NewServer(shim)
+	defer shimServer.Close()
+
+	// 1. Open a websocket connection via the shim.
+	openURL := shimServer.URL + shimPath + "open"
+	resp, err := http.Post(openURL, "text/plain", strings.NewReader(backendURL.String()))
+	if err != nil {
+		t.Fatalf("Failed to open shim connection: %v", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("Failed to open shim connection, status: %d", resp.StatusCode)
+	}
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("Failed to read open response body: %v", err)
+	}
+	resp.Body.Close()
+	var openResp sessionMessage
+	if err := json.Unmarshal(body, &openResp); err != nil {
+		t.Fatalf("Failed to unmarshal open response: %v", err)
+	}
+	sessionID := openResp.ID
+	if sessionID == "" {
+		t.Fatal("No sessionID in open response")
+	}
+
+	// 2. Poll repeatedly without any messages being sent.
+	pollURL := shimServer.URL + shimPath + "poll"
+	pollReq := fmt.Sprintf(`{"id": %q}`, sessionID)
+	timeout := time.Now().Add(idleTimeout + 1*time.Second)
+	for time.Now().Before(timeout) {
+		resp, err := http.Post(pollURL, "application/json", strings.NewReader(pollReq))
+		if err != nil {
+			t.Fatalf("Failed to poll shim connection: %v", err)
+		}
+		resp.Body.Close()
+		if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusRequestTimeout {
+			t.Fatalf("Unexpected status code during polling: %d", resp.StatusCode)
+		}
+		// Sleep for half of read timeout to simulate polling faster than idle timeout.
+		time.Sleep(500 * time.Millisecond)
+	}
+
+	// 3. After idleTimeout + 1s, one more poll should succeed.
+	resp, err = http.Post(pollURL, "application/json", strings.NewReader(pollReq))
+	if err != nil {
+		t.Fatalf("Failed to poll shim connection: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusRequestTimeout {
+		t.Errorf("Polling after idle timeout got status %d, want %d", resp.StatusCode, http.StatusRequestTimeout)
+	}
+
+	// 4. Close connection.
+	closeURL := shimServer.URL + shimPath + "close"
+	closeReq := fmt.Sprintf(`{"id": %q}`, sessionID)
+	resp, err = http.Post(closeURL, "application/json", strings.NewReader(closeReq))
+	if err != nil {
+		t.Fatalf("Failed to close shim connection: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("Close shim connection got status %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+}

testing/websockets/main.go

@@ -31,6 +31,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"time"
 
 	"github.com/google/inverting-proxy/agent/metrics"
 	"github.com/google/inverting-proxy/agent/websockets"
@@ -48,6 +49,7 @@ var (
 	monitoringEndpoint       = flag.String("monitoring-endpoint", "staging-monitoring.sandbox.googleapis.com:443", "The endpoint to which to write metrics. Eg: monitoring.googleapis.com corresponds to Cloud Monarch.")
 	monitoringResourceType   = flag.String("monitoring-resource-type", "gce_instance", "The monitoring resource type. Eg: gce_instance")
 	monitoringResourceLabels = flag.String("monitoring-resource-labels", "instance-id=fake-instance-id,instance-zone=us-west1-a", "Comma separated key value pairs for the purpose of monitoring configuration. Eg: 'instance-id=my-instance-id,instance-zone=us-west1-a")
+	websocketShimTimeout     = flag.Duration("websocket-shim-timeout", 60*time.Minute, "Timeout for websocket shim connections to expire due to inactivity.")
 )
 
 func main() {
@@ -69,7 +71,7 @@ func main() {
 	}
 
 	backendProxy := httputil.NewSingleHostReverseProxy(backendURL)
-	shimmingProxy, err := websockets.Proxy(context.Background(), backendProxy, backendURL.Host, *shimPath, true, *enableWebsocketInjection, func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler { return h }, metricHandler)
+	shimmingProxy, err := websockets.Proxy(context.Background(), backendProxy, backendURL.Host, *shimPath, true, *enableWebsocketInjection, func(h http.Handler, metricHandler *metrics.MetricHandler) http.Handler { return h }, metricHandler, *websocketShimTimeout)
 	if err != nil {
 		log.Fatalf("Failure starting the websocket-shimming proxy: %v", err)
 	}