fix(envd): stop memory exhaustion when client disconnects from streaming process

packages/envd/internal/services/process/handler/handler.go

@@ -51,11 +51,11 @@ type Handler struct {
 
 	cancel context.CancelFunc
 
-	outCtx    context.Context //nolint:containedctx // todo: refactor so this can be removed
-	outCancel context.CancelFunc
+	readersDone chan struct{} // closed when stdout/stderr reader goroutines have exited
 
-	stdinMu sync.Mutex
-	stdin   io.WriteCloser
+	stdinMu  sync.Mutex
+	stdin    io.WriteCloser
+	pipeRead []*os.File // read-ends of stdout/stderr
 
 	DataEvent *MultiplexedChannel[rpc.ProcessEvent_Data]
 	EndEvent  *MultiplexedChannel[rpc.ProcessEvent_End]
@@ -172,20 +172,15 @@ func New(
 
 	var outWg sync.WaitGroup
 
-	// Create a context for waiting for and cancelling output pipes.
-	// Cancellation of the process via timeout will propagate and cancel this context too.
-	outCtx, outCancel := context.WithCancel(ctx)
-
 	h := &Handler{
-		Config:    req.GetProcess(),
-		cmd:       cmd,
-		Tag:       req.Tag,
-		DataEvent: outMultiplex,
-		cancel:    cancel,
-		outCtx:    outCtx,
-		outCancel: outCancel,
-		EndEvent:  NewMultiplexedChannel[rpc.ProcessEvent_End](0),
-		logger:    logger,
+		Config:      req.GetProcess(),
+		cmd:         cmd,
+		Tag:         req.Tag,
+		DataEvent:   outMultiplex,
+		cancel:      cancel,
+		readersDone: make(chan struct{}),
+		EndEvent:    NewMultiplexedChannel[rpc.ProcessEvent_End](1),
+		logger:      logger,
 	}
 
 	if req.GetPty() != nil {
@@ -206,16 +201,18 @@ func New(
 				n, readErr := tty.Read(buf)
 
 				if n > 0 {
-					outMultiplex.Source <- rpc.ProcessEvent_Data{
+					event := rpc.ProcessEvent_Data{
 						Data: &rpc.ProcessEvent_DataEvent{
 							Output: &rpc.ProcessEvent_DataEvent_Pty{
 								Pty: buf[:n],
 							},
 						},
 					}
+
+					outMultiplex.Source <- event
 				}
 
-				if errors.Is(readErr, io.EOF) {
+				if errors.Is(readErr, io.EOF) || os.IsTimeout(readErr) {
 					break
 				}
 
@@ -229,12 +226,17 @@ func New(
 
 		h.tty = tty
 	} else {
-		stdout, err := cmd.StdoutPipe()
+		stdoutR, stdoutW, err := os.Pipe()
 		if err != nil {
 			return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("error creating stdout pipe for command '%s': %w", userCmd, err))
 		}
 
+		cmd.Stdout = stdoutW
+		stdout := stdoutR
+
 		outWg.Go(func() {
+			defer stdout.Close()
+
 			stdoutLogs := make(chan []byte, outputBufferSize)
 			defer close(stdoutLogs)
 
@@ -248,18 +250,20 @@ func New(
 				n, readErr := stdout.Read(buf)
 
 				if n > 0 {
-					outMultiplex.Source <- rpc.ProcessEvent_Data{
+					event := rpc.ProcessEvent_Data{
 						Data: &rpc.ProcessEvent_DataEvent{
 							Output: &rpc.ProcessEvent_DataEvent_Stdout{
 								Stdout: buf[:n],
 							},
 						},
 					}
 
+					outMultiplex.Source <- event
+
 					stdoutLogs <- buf[:n]
 				}
 
-				if errors.Is(readErr, io.EOF) {
+				if errors.Is(readErr, io.EOF) || os.IsTimeout(readErr) {
 					break
 				}
 
@@ -271,12 +275,19 @@ func New(
 			}
 		})
 
-		stderr, err := cmd.StderrPipe()
+		stderrR, stderrW, err := os.Pipe()
 		if err != nil {
 			return nil, connect.NewError(connect.CodeInternal, fmt.Errorf("error creating stderr pipe for command '%s': %w", userCmd, err))
 		}
 
+		cmd.Stderr = stderrW
+		stderr := stderrR
+
+		h.pipeRead = []*os.File{stdoutR, stderrR}
+
 		outWg.Go(func() {
+			defer stderr.Close()
+
 			stderrLogs := make(chan []byte, outputBufferSize)
 			defer close(stderrLogs)
 
@@ -290,18 +301,20 @@ func New(
 				n, readErr := stderr.Read(buf)
 
 				if n > 0 {
-					outMultiplex.Source <- rpc.ProcessEvent_Data{
+					event := rpc.ProcessEvent_Data{
 						Data: &rpc.ProcessEvent_DataEvent{
 							Output: &rpc.ProcessEvent_DataEvent_Stderr{
 								Stderr: buf[:n],
 							},
 						},
 					}
 
+					outMultiplex.Source <- event
+
 					stderrLogs <- buf[:n]
 				}
 
-				if errors.Is(readErr, io.EOF) {
+				if errors.Is(readErr, io.EOF) || os.IsTimeout(readErr) {
 					break
 				}
 
@@ -328,9 +341,9 @@ func New(
 	go func() {
 		outWg.Wait()
 
-		close(outMultiplex.Source)
+		close(h.readersDone)
 
-		outCancel()
+		outMultiplex.CloseSource()
 	}()
 
 	return h, nil
@@ -349,10 +362,6 @@ func (p *Handler) SendSignal(signal syscall.Signal) error {
 		return errors.New("process not started")
 	}
 
-	if signal == syscall.SIGKILL || signal == syscall.SIGTERM {
-		p.outCancel()
-	}
-
 	return p.cmd.Process.Signal(signal)
 }
 
@@ -426,6 +435,16 @@ func (p *Handler) Start(requestTimeout time.Duration) (uint32, error) {
 		if err != nil {
 			return 0, fmt.Errorf("error starting process '%s': %w", p.userCommand(), err)
 		}
+
+		// Close parent's copy of the write-ends so readers see EOF
+		// when the child (and any orphan grandchildren) exit.
+		if p.cmd.Stdout != nil {
+			p.cmd.Stdout.(*os.File).Close()
+		}
+
+		if p.cmd.Stderr != nil {
+			p.cmd.Stderr.(*os.File).Close()
+		}
 	}
 
 	p.logger.
@@ -440,11 +459,30 @@ func (p *Handler) Start(requestTimeout time.Duration) (uint32, error) {
 }
 
 func (p *Handler) Wait() {
-	// Wait for the output pipes to be closed or cancelled.
-	<-p.outCtx.Done()
-
+	// Reap the child. With manual os.Pipe, cmd.Wait does not
+	// close our read-ends, so readers can finish draining.
 	err := p.cmd.Wait()
 
+	// Disable back-pressure so any reader stuck on a full Source
+	// send unblocks, loops back, and sees EOF.
+	p.DataEvent.Drain()
+
+	// Set a read deadline on the pipe/pty read-ends so readers drain
+	// any buffered data (reads with available data return instantly)
+	// and then exit cleanly instead of blocking forever when an
+	// orphan grandchild holds the write-end open.
+	deadline := time.Now().Add(1 * time.Second)
+
+	for _, f := range p.pipeRead {
+		f.SetReadDeadline(deadline)
+	}
+
+	if p.tty != nil {
+		p.tty.SetReadDeadline(deadline)
+	}
+
+	<-p.readersDone
+
 	p.tty.Close()
 
 	var errMsg *string
@@ -466,6 +504,7 @@ func (p *Handler) Wait() {
 	}
 
 	p.EndEvent.Source <- event
+	p.EndEvent.CloseSource()
 
 	p.logger.
 		Info().