Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions e2e/framework/retry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
23 changes: 13 additions & 10 deletions pkg/tunnel/container.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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",
Expand Down
56 changes: 55 additions & 1 deletion pkg/tunnel/container_test.go
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down Expand Up @@ -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)
}
Loading