mirror of
https://github.com/cirruslabs/orchard.git
synced 2026-10-22 06:18:16 +02:00
Supporting reconnecting to /exec socket (#434)
This commit is contained in:
Fedor Kororkov
@@ -0,0 +1,320 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cirruslabs/orchard/internal/execstream"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fakeExec struct {
|
||||
stdin io.WriteCloser
|
||||
run func(context.Context, string, chan<- *execstream.Frame) error
|
||||
closeCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Stdin() io.WriteCloser {
|
||||
return exec.stdin
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Run(
|
||||
ctx context.Context,
|
||||
command string,
|
||||
outgoingFrames chan<- *execstream.Frame,
|
||||
) error {
|
||||
if exec.run != nil {
|
||||
return exec.run(ctx, command, outgoingFrames)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (exec *fakeExec) Close() error {
|
||||
exec.closeCalls.Add(1)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func newManualExecSessionForTest(
|
||||
key execSessionKey,
|
||||
registry *execSessionRegistry,
|
||||
) *execSession {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
return &execSession{
|
||||
key: key,
|
||||
command: "echo test",
|
||||
exec: &fakeExec{},
|
||||
registry: registry,
|
||||
exitTTL: time.Minute,
|
||||
policy: reconnectableExecSessionPolicy,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
subscribers: map[*execSessionSubscriber]struct{}{},
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
|
||||
createStarted := make(chan struct{})
|
||||
releaseCreate := make(chan struct{})
|
||||
var createCalls atomic.Int32
|
||||
|
||||
create := func() (*execSession, error) {
|
||||
createCalls.Add(1)
|
||||
close(createStarted)
|
||||
<-releaseCreate
|
||||
|
||||
return newManualExecSessionForTest(key, registry), nil
|
||||
}
|
||||
|
||||
firstDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(firstDone)
|
||||
_, _, err := registry.getOrCreate(context.Background(), key, create)
|
||||
require.NoError(t, err)
|
||||
}()
|
||||
|
||||
<-createStarted
|
||||
|
||||
secondDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(secondDone)
|
||||
_, created, err := registry.getOrCreate(context.Background(), key, create)
|
||||
require.NoError(t, err)
|
||||
require.False(t, created)
|
||||
}()
|
||||
|
||||
close(releaseCreate)
|
||||
|
||||
<-firstDone
|
||||
<-secondDone
|
||||
require.EqualValues(t, 1, createCalls.Load())
|
||||
}
|
||||
|
||||
func TestExecSessionStartRunsCommandOnlyOnce(t *testing.T) {
|
||||
var runCalls atomic.Int32
|
||||
runStarted := make(chan struct{})
|
||||
|
||||
session := newExecSession(
|
||||
execSessionKey{vmName: "vm", sessionID: "session"},
|
||||
"echo test",
|
||||
&fakeExec{
|
||||
run: func(ctx context.Context, _ string, _ chan<- *execstream.Frame) error {
|
||||
runCalls.Add(1)
|
||||
close(runStarted)
|
||||
<-ctx.Done()
|
||||
|
||||
return ctx.Err()
|
||||
},
|
||||
},
|
||||
nil,
|
||||
nil,
|
||||
time.Minute,
|
||||
reconnectableExecSessionPolicy,
|
||||
)
|
||||
defer session.close()
|
||||
|
||||
session.start()
|
||||
session.start()
|
||||
|
||||
<-runStarted
|
||||
require.EqualValues(t, 1, runCalls.Load())
|
||||
}
|
||||
|
||||
func TestExecSessionHistoryReplayAndAck(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStderr, Data: []byte("err")})
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeExit,
|
||||
Exit: &execstream.Exit{Code: 7},
|
||||
})
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.sendHistory(subscriber, 0)
|
||||
|
||||
require.Equal(t, execstream.FrameTypeStdout, (<-subscriber.frames).Type)
|
||||
require.Equal(t, execstream.FrameTypeStderr, (<-subscriber.frames).Type)
|
||||
require.Equal(t, execstream.FrameTypeExit, (<-subscriber.frames).Type)
|
||||
noMoreHistory := <-subscriber.frames
|
||||
require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type)
|
||||
require.EqualValues(t, 3, noMoreHistory.Watermark)
|
||||
|
||||
session.ack(2)
|
||||
require.Len(t, session.replay.frames, 1)
|
||||
require.EqualValues(t, 3, session.replay.frames[0].frame.Watermark)
|
||||
}
|
||||
|
||||
func TestExecSessionHistoryReplayStreamsPastSubscriberBuffer(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
|
||||
const frameCount = 256
|
||||
for i := 0; i < frameCount; i++ {
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeStdout,
|
||||
Data: []byte{byte(i)},
|
||||
})
|
||||
}
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
session.sendHistory(subscriber, 0)
|
||||
}()
|
||||
|
||||
for i := 1; i <= frameCount; i++ {
|
||||
frame := <-subscriber.frames
|
||||
require.Equal(t, execstream.FrameTypeStdout, frame.Type)
|
||||
require.EqualValues(t, i, frame.Watermark)
|
||||
}
|
||||
|
||||
noMoreHistory := <-subscriber.frames
|
||||
require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type)
|
||||
require.EqualValues(t, frameCount, noMoreHistory.Watermark)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-done:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestExecSessionDetachKeepsProcessAlive(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
exec := session.exec.(*fakeExec)
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.detach(subscriber)
|
||||
|
||||
require.False(t, session.closed)
|
||||
require.EqualValues(t, 0, exec.closeCalls.Load())
|
||||
}
|
||||
|
||||
func TestLegacyExecSessionDetachStopsProcess(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
session.policy = legacyExecSessionPolicy
|
||||
exec := session.exec.(*fakeExec)
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.detach(subscriber)
|
||||
|
||||
require.True(t, session.closed)
|
||||
require.EqualValues(t, 1, exec.closeCalls.Load())
|
||||
}
|
||||
|
||||
func TestLegacyExecSessionDoesNotRetainReplayHistory(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
session.policy = legacyExecSessionPolicy
|
||||
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
|
||||
|
||||
require.Empty(t, session.replay.frames)
|
||||
require.Zero(t, session.replay.nextWatermark)
|
||||
}
|
||||
|
||||
func TestExecSessionCloseIfUnusedClosesIdleSession(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
session := newManualExecSessionForTest(key, registry)
|
||||
exec := session.exec.(*fakeExec)
|
||||
registry.sessions[key] = session
|
||||
|
||||
session.closeIfUnused()
|
||||
|
||||
require.True(t, session.closed)
|
||||
require.EqualValues(t, 1, exec.closeCalls.Load())
|
||||
}
|
||||
|
||||
func TestExecSessionCloseIfUnusedKeepsAttachedSession(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
|
||||
_, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.closeIfUnused()
|
||||
|
||||
require.False(t, session.closed)
|
||||
}
|
||||
|
||||
func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
session := newManualExecSessionForTest(key, registry)
|
||||
exec := session.exec.(*fakeExec)
|
||||
registry.sessions[key] = session
|
||||
|
||||
session.close()
|
||||
|
||||
require.True(t, session.closed)
|
||||
require.EqualValues(t, 1, exec.closeCalls.Load())
|
||||
_, ok := registry.get(key)
|
||||
require.False(t, ok)
|
||||
}
|
||||
|
||||
func TestExecSessionFinishedEntryExpiresAfterTTL(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
key := execSessionKey{vmName: "vm", sessionID: "session"}
|
||||
session := newManualExecSessionForTest(key, registry)
|
||||
session.exitTTL = 10 * time.Millisecond
|
||||
registry.sessions[key] = session
|
||||
|
||||
session.markFinished()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
_, ok := registry.get(key)
|
||||
|
||||
return !ok
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestExecSessionFinishKeepsReconnectableSubscriberOpen(t *testing.T) {
|
||||
registry := newExecSessionRegistry()
|
||||
session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
|
||||
|
||||
subscriber, err := session.attach()
|
||||
require.NoError(t, err)
|
||||
|
||||
session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeExit,
|
||||
Exit: &execstream.Exit{Code: 0},
|
||||
})
|
||||
session.markFinished()
|
||||
|
||||
require.Equal(t, execstream.FrameTypeStdout, (<-subscriber.frames).Type)
|
||||
require.Equal(t, execstream.FrameTypeExit, (<-subscriber.frames).Type)
|
||||
|
||||
session.sendHistory(subscriber, 0)
|
||||
|
||||
noMoreHistory, ok := <-subscriber.frames
|
||||
require.True(t, ok)
|
||||
require.Equal(t, execstream.FrameTypeNoMoreHistory, noMoreHistory.Type)
|
||||
require.EqualValues(t, 2, noMoreHistory.Watermark)
|
||||
}
|
||||
Reference in New Issue
Block a user