refactor buffer size

pkg/abstractions/shell/http.go

@@ -81,21 +81,37 @@ func (g *shellGroup) ShellConnect(ctx echo.Context) error {
 	// TODO: confirm disconnects happen when python client exits
 	defer log.Println("disconnected")
 
+	// Create a context that will be canceled when the client disconnects
+	clientCtx, clientCancel := context.WithCancel(ctx.Request().Context())
+	defer clientCancel()
+
 	var wg sync.WaitGroup
 	wg.Add(2)
 
 	// Copy from client -> container
 	go func() {
 		defer wg.Done()
-		buf := make([]byte, 32*1024) // 32KB buffer
-		io.CopyBuffer(containerConn, conn, buf)
+		buf := make([]byte, shellProxyBufferSizeKb)
+		select {
+		case <-clientCtx.Done():
+			log.Println("client disconnected")
+			return
+		default:
+			io.CopyBuffer(containerConn, conn, buf)
+		}
 	}()
 
 	// Copy from container -> client
 	go func() {
 		defer wg.Done()
-		buf := make([]byte, 32*1024) // 32KB buffer
-		io.CopyBuffer(conn, containerConn, buf)
+		buf := make([]byte, shellProxyBufferSizeKb)
+		select {
+		case <-clientCtx.Done():
+			log.Println("client disconnected")
+			return
+		default:
+			io.CopyBuffer(conn, containerConn, buf)
+		}
 	}()
 
 	wg.Wait()

pkg/abstractions/shell/shell.go

@@ -25,6 +25,7 @@ import (
 const (
 	shellRoutePrefix              string        = "/shell"
 	shellContainerPrefix          string        = "shell"
+	shellProxyBufferSizeKb        int           = 32 * 1024
 	defaultContainerCpu           int64         = 100
 	defaultContainerMemory        int64         = 128
 	containerDialTimeoutDurationS time.Duration = 300 * time.Second