diff --git a/e2e/framework/retry_test.go b/e2e/framework/retry_test.go index 0ccc08961e..2a19f98869 100644 --- a/e2e/framework/retry_test.go +++ b/e2e/framework/retry_test.go @@ -422,6 +422,25 @@ func TestExecWithSSHRetry_RetryHandshakeStallThenSuccess(t *testing.T) { assert.Equal(t, "ok", out) } +func TestExecWithSSHRetry_RetryStartupSilenceEnvelopeThenSuccess(t *testing.T) { + withFastBackoffs(t) + calls := 0 + envelope := `{"kind":"error","outcome":"error","code":"UNKNOWN","message":` + + `"container exec session produced no output for 30s: remote process never started"}` + out, err := execWithSSHRetry(context.Background(), "ws", + func(context.Context) (string, string, error) { + calls++ + if calls == 1 { + return "", envelope, transientExitErr(t) + } + return "ok", "", nil + }, + ) + require.NoError(t, err) + assert.Equal(t, 2, calls) + assert.Equal(t, "ok", out) +} + func TestExecWithSSHRetry_RetryExhausted(t *testing.T) { withFastBackoffs(t) calls := 0 diff --git a/pkg/tunnel/container.go b/pkg/tunnel/container.go index 7948d3f8ee..81b4d4b049 100644 --- a/pkg/tunnel/container.go +++ b/pkg/tunnel/container.go @@ -118,11 +118,8 @@ func (c *ContainerTunnel) runHostTunnel( stdinReader io.Reader, stdoutWriter io.Writer, timeout time.Duration, ) error { - writer, done := log.PipeJSONStreamWithFallback(log.PassthroughWriter()) - defer func() { - _ = writer.Close() - <-done - }() + writer, closeAndDrain := internalProcessStderr(log.PassthroughWriter()) + defer closeAndDrain() defer log.Debugf("Tunnel to host closed") command := fmt.Sprintf("'%s' internal ssh-server --stdio", c.client.AgentPath()) @@ -150,6 +147,15 @@ func (c *ContainerTunnel) runHostTunnel( }) } +func internalProcessStderr(fallback io.WriteCloser) (io.WriteCloser, func()) { + writer, done := log.PipeJSONStreamWithFallback(fallback) + return writer, func() { + _ = writer.Close() + <-done + _ = fallback.Close() + } +} + // updateConfig is called periodically to keep the workspace agent config up to date. func (c *ContainerTunnel) updateConfig(ctx context.Context, sshClient *ssh.Client) { for { @@ -209,11 +215,8 @@ func (c *ContainerTunnel) runInContainer( return err } - writer, writerDone := log.PipeJSONStream() - defer func() { - _ = writer.Close() - <-writerDone - }() + writer, closeAndDrain := internalProcessStderr(log.PassthroughWriter()) + defer closeAndDrain() command := fmt.Sprintf( "%q internal agent container-tunnel --workspace-info %q", diff --git a/pkg/tunnel/container_test.go b/pkg/tunnel/container_test.go index aa83a6c343..947819da2a 100644 --- a/pkg/tunnel/container_test.go +++ b/pkg/tunnel/container_test.go @@ -1,11 +1,16 @@ package tunnel import ( + "bytes" "context" + "io" "sync" "sync/atomic" "testing" "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) type fakeCloser struct { @@ -66,8 +71,57 @@ func TestStopUpdateThenCloseWaitsForGoroutine(t *testing.T) { case <-time.After(2 * time.Second): t.Fatal("stopUpdateThenClose did not return after goroutine finished") } - if !closer.closed.Load() { t.Fatal("expected sshClient to be closed") } } + +type closeBuffer struct { + bytes.Buffer + closed bool +} + +func (b *closeBuffer) Close() error { + b.closed = true + return nil +} + +func TestInternalProcessStderrPreservesNonLogOutput(t *testing.T) { + fallback := &closeBuffer{} + writer, closeAndDrain := internalProcessStderr(fallback) + + errorEnvelope := `{"kind":"error","outcome":"error","code":"UNKNOWN","message":` + + `"container exec session produced no output for 30s: remote process never started"}` + _, err := io.WriteString(writer, + `{"level":"debug","msg":"structured-log-marker"}`+"\n"+ + errorEnvelope+"\n"+ + "plain subprocess failure\n") + require.NoError(t, err) + closeAndDrain() + + got := fallback.String() + assert.NotContains(t, got, "structured-log-marker") + assert.Contains(t, got, errorEnvelope) + assert.Contains(t, got, "plain subprocess failure\n") + assert.Equal( + t, + 1, + bytes.Count([]byte(got), []byte("container exec session produced no output")), + ) + assert.True(t, fallback.closed) +} + +func TestInternalProcessStderrDrainsFinalUnterminatedLine(t *testing.T) { + fallback := &closeBuffer{} + writer, closeAndDrain := internalProcessStderr(fallback) + + _, err := io.WriteString(writer, + `{"kind":"error","message":"container exec session produced no output"}`) + require.NoError(t, err) + closeAndDrain() + + assert.Equal(t, + `{"kind":"error","message":"container exec session produced no output"}`+"\n", + fallback.String()) + assert.True(t, fallback.closed) +}