[rpc] Preserve background processes after successful exec RPCs (#64)

2026-09-29 19:31:10 +02:00 · 2026-08-27 18:05:37 -07:00
parent 94a27a6724
commit cb39b12dcc
2 changed files with 137 additions and 1 deletions
@@ -19,6 +19,8 @@ import (
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/emptypb"
 )

@@ -199,7 +201,7 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse])
 					return
 				}

-				if !errors.Is(err, context.Canceled) {
+				if !errors.Is(err, context.Canceled) && status.Code(err) != codes.Canceled {
 					reportClientError(err)
 				}

@@ -0,0 +1,134 @@
+//nolint:testpackage
+package rpc
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+func TestExecGRPCBackgroundProcessLifetime(t *testing.T) {
+	for _, mode := range []string{"normal exit", "cancel after exit", "cancel while running"} {
+		t.Run(mode, func(t *testing.T) {
+			client := NewAgentClient(newExecGRPCTestConn(t))
+			ctx, cancel := context.WithTimeout(t.Context(), execTestTimeout)
+			defer cancel()
+
+			stream, err := client.Exec(ctx)
+			require.NoError(t, err)
+
+			releasePath := filepath.Join(t.TempDir(), "release")
+			require.NoError(t, stream.Send(&ExecRequest{Type: &ExecRequest_Command_{Command: &ExecRequest_Command{
+				Name: execTestShell,
+				Args: []string{"-c", `nohup sleep 30 >/dev/null 2>&1 </dev/null &
+echo "$!"
+while [ ! -f "$RELEASE_FILE" ]; do sleep 0.01; done`},
+				Env: map[string]string{"RELEASE_FILE": releasePath},
+			}}}))
+
+			var output strings.Builder
+			for !strings.Contains(output.String(), "\n") {
+				response, err := stream.Recv()
+				require.NoError(t, err)
+				require.Nil(t, response.GetExit())
+				output.Write(response.GetStandardOutput().GetData())
+			}
+			childPID, err := strconv.Atoi(strings.TrimSpace(output.String()))
+			require.NoError(t, err)
+			require.Greater(t, childPID, 1)
+			t.Cleanup(func() {
+				if execTestProcessRunning(t, childPID) {
+					require.NoError(t, syscall.Kill(childPID, syscall.SIGKILL))
+				}
+				require.Eventually(t, func() bool {
+					return !execTestProcessRunning(t, childPID)
+				}, execTestTimeout, 10*time.Millisecond)
+			})
+			require.True(t, execTestProcessRunning(t, childPID), "child must start before launcher completes")
+
+			if mode == "cancel while running" {
+				cancel()
+				_, err := stream.Recv()
+				require.Equal(t, codes.Canceled, status.Code(err))
+				require.Eventually(t, func() bool {
+					return !execTestProcessRunning(t, childPID)
+				}, execTestTimeout, 10*time.Millisecond, "running command cancellation must kill its child")
+
+				return
+			}
+
+			require.NoError(t, os.WriteFile(releasePath, nil, 0o600))
+			response, err := stream.Recv()
+			require.NoError(t, err)
+			require.NotNil(t, response.GetExit())
+			require.Zero(t, response.GetExit().GetCode())
+			if mode == "cancel after exit" {
+				cancel()
+			} else {
+				_, err = stream.Recv()
+				require.ErrorIs(t, err, io.EOF)
+			}
+
+			require.Never(t, func() bool {
+				return !execTestProcessRunning(t, childPID)
+			}, 250*time.Millisecond, 10*time.Millisecond, "RPC shutdown must preserve background children")
+		})
+	}
+}
+
+func newExecGRPCTestConn(t *testing.T) *grpc.ClientConn {
+	listener := bufconn.Listen(1024 * 1024)
+	agent, err := New(listener)
+	require.NoError(t, err)
+
+	serveResult := make(chan error, 1)
+	go func() { serveResult <- agent.grpcServer.Serve(listener) }()
+	t.Cleanup(func() {
+		agent.grpcServer.Stop()
+		require.NoError(t, <-serveResult)
+	})
+
+	conn, err := grpc.NewClient("passthrough:///exec-test",
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) { return listener.DialContext(ctx) }),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, conn.Close()) })
+
+	return conn
+}
+
+func execTestProcessRunning(t *testing.T, pid int) bool {
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(t.Context()), execTestTimeout)
+	defer cancel()
+
+	//nolint:gosec
+	output, err := exec.CommandContext(ctx, "/bin/ps", "-p", strconv.Itoa(pid), "-o", "stat=").Output()
+	if err != nil {
+		var exitError *exec.ExitError
+		if errors.As(err, &exitError) && exitError.ExitCode() == 1 {
+			return false
+		}
+		require.NoError(t, err)
+	}
+
+	state := strings.TrimSpace(string(output))
+
+	return state != "" && !strings.HasPrefix(state, "Z")
+}