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,692 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/cirruslabs/orchard/internal/execstream"
|
||||
)
|
||||
|
||||
const execSessionReplayBufferBytes = 4 * 1024 * 1024
|
||||
|
||||
type execSessionPolicy struct {
|
||||
closeOnDetach bool
|
||||
retainAfterExit bool
|
||||
replayEnabled bool
|
||||
}
|
||||
|
||||
var (
|
||||
legacyExecSessionPolicy = execSessionPolicy{
|
||||
closeOnDetach: true,
|
||||
}
|
||||
reconnectableExecSessionPolicy = execSessionPolicy{
|
||||
retainAfterExit: true,
|
||||
replayEnabled: true,
|
||||
}
|
||||
)
|
||||
|
||||
type sshExecRunner interface {
|
||||
Stdin() io.WriteCloser
|
||||
Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error
|
||||
Close() error
|
||||
}
|
||||
|
||||
type execSessionKey struct {
|
||||
vmName string
|
||||
sessionID string
|
||||
}
|
||||
|
||||
type execSessionCreation struct {
|
||||
done chan struct{}
|
||||
session *execSession
|
||||
err error
|
||||
}
|
||||
|
||||
type execSessionRegistry struct {
|
||||
mu sync.Mutex
|
||||
sessions map[execSessionKey]*execSession
|
||||
creating map[execSessionKey]*execSessionCreation
|
||||
}
|
||||
|
||||
func newExecSessionRegistry() *execSessionRegistry {
|
||||
return &execSessionRegistry{
|
||||
sessions: map[execSessionKey]*execSession{},
|
||||
creating: map[execSessionKey]*execSessionCreation{},
|
||||
}
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) get(key execSessionKey) (*execSession, bool) {
|
||||
registry.mu.Lock()
|
||||
defer registry.mu.Unlock()
|
||||
|
||||
session, ok := registry.sessions[key]
|
||||
|
||||
return session, ok
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) getOrCreate(
|
||||
ctx context.Context,
|
||||
key execSessionKey,
|
||||
create func() (*execSession, error),
|
||||
) (*execSession, bool, error) {
|
||||
registry.mu.Lock()
|
||||
|
||||
if session, ok := registry.sessions[key]; ok {
|
||||
registry.mu.Unlock()
|
||||
|
||||
return session, false, nil
|
||||
}
|
||||
|
||||
if creation, ok := registry.creating[key]; ok {
|
||||
registry.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, false, ctx.Err()
|
||||
case <-creation.done:
|
||||
if creation.err != nil {
|
||||
return nil, false, creation.err
|
||||
}
|
||||
|
||||
return creation.session, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
creation := &execSessionCreation{done: make(chan struct{})}
|
||||
registry.creating[key] = creation
|
||||
registry.mu.Unlock()
|
||||
|
||||
session, err := create()
|
||||
|
||||
registry.mu.Lock()
|
||||
delete(registry.creating, key)
|
||||
if err == nil {
|
||||
registry.sessions[key] = session
|
||||
}
|
||||
creation.session = session
|
||||
creation.err = err
|
||||
close(creation.done)
|
||||
registry.mu.Unlock()
|
||||
|
||||
return session, true, err
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) remove(key execSessionKey, expected *execSession) {
|
||||
registry.mu.Lock()
|
||||
defer registry.mu.Unlock()
|
||||
|
||||
if registry.sessions[key] == expected {
|
||||
delete(registry.sessions, key)
|
||||
}
|
||||
}
|
||||
|
||||
func (registry *execSessionRegistry) closeAll() {
|
||||
registry.mu.Lock()
|
||||
sessions := make([]*execSession, 0, len(registry.sessions))
|
||||
for _, session := range registry.sessions {
|
||||
sessions = append(sessions, session)
|
||||
}
|
||||
registry.mu.Unlock()
|
||||
|
||||
for _, session := range sessions {
|
||||
session.close()
|
||||
}
|
||||
}
|
||||
|
||||
type execReplayFrame struct {
|
||||
frame *execstream.Frame
|
||||
size int
|
||||
}
|
||||
|
||||
type execReplayBuffer struct {
|
||||
frames []execReplayFrame
|
||||
bufferBytes int
|
||||
nextWatermark uint64
|
||||
ackedWatermark uint64
|
||||
}
|
||||
|
||||
func (buffer *execReplayBuffer) append(frame *execstream.Frame) *execstream.Frame {
|
||||
frame = cloneExecFrame(frame)
|
||||
buffer.nextWatermark++
|
||||
frame.Watermark = buffer.nextWatermark
|
||||
|
||||
frameSize := execFrameSize(frame)
|
||||
buffer.frames = append(buffer.frames, execReplayFrame{
|
||||
frame: frame,
|
||||
size: frameSize,
|
||||
})
|
||||
buffer.bufferBytes += frameSize
|
||||
buffer.trimAcknowledged()
|
||||
buffer.trimToLimit()
|
||||
|
||||
return frame
|
||||
}
|
||||
|
||||
func (buffer *execReplayBuffer) ack(watermark uint64) {
|
||||
if watermark <= buffer.ackedWatermark {
|
||||
return
|
||||
}
|
||||
|
||||
buffer.ackedWatermark = watermark
|
||||
buffer.trimAcknowledged()
|
||||
}
|
||||
|
||||
func (buffer *execReplayBuffer) replayAfter(
|
||||
watermark uint64,
|
||||
frames []*execstream.Frame,
|
||||
) []*execstream.Frame {
|
||||
for _, record := range buffer.frames {
|
||||
if record.frame.Watermark <= watermark {
|
||||
continue
|
||||
}
|
||||
|
||||
frames = append(frames, record.frame)
|
||||
}
|
||||
|
||||
return frames
|
||||
}
|
||||
|
||||
func (buffer *execReplayBuffer) trimAcknowledged() {
|
||||
for len(buffer.frames) > 0 && buffer.frames[0].frame.Watermark <= buffer.ackedWatermark {
|
||||
buffer.bufferBytes -= buffer.frames[0].size
|
||||
buffer.frames = buffer.frames[1:]
|
||||
}
|
||||
}
|
||||
|
||||
func (buffer *execReplayBuffer) trimToLimit() {
|
||||
for buffer.bufferBytes > execSessionReplayBufferBytes && len(buffer.frames) > 0 {
|
||||
buffer.bufferBytes -= buffer.frames[0].size
|
||||
buffer.frames = buffer.frames[1:]
|
||||
}
|
||||
}
|
||||
|
||||
type execSessionSubscriber struct {
|
||||
frames chan *execstream.Frame
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
sendMu sync.Mutex
|
||||
sentWatermark uint64
|
||||
}
|
||||
|
||||
func newExecSessionSubscriber() *execSessionSubscriber {
|
||||
return &execSessionSubscriber{
|
||||
frames: make(chan *execstream.Frame, 128),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) enqueue(frame *execstream.Frame) bool {
|
||||
subscriber.sendMu.Lock()
|
||||
defer subscriber.sendMu.Unlock()
|
||||
|
||||
if subscriber.alreadySentLocked(frame) {
|
||||
return true
|
||||
}
|
||||
|
||||
select {
|
||||
case <-subscriber.closed:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
|
||||
select {
|
||||
case subscriber.frames <- subscriber.markSentLocked(frame):
|
||||
return true
|
||||
case <-subscriber.closed:
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) sendHistory(frames []*execstream.Frame) bool {
|
||||
for _, frame := range frames {
|
||||
if !subscriber.sendLocked(frame) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) sendLocked(frame *execstream.Frame) bool {
|
||||
if subscriber.alreadySentLocked(frame) {
|
||||
return true
|
||||
}
|
||||
|
||||
select {
|
||||
case <-subscriber.closed:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
|
||||
select {
|
||||
case subscriber.frames <- subscriber.markSentLocked(frame):
|
||||
return true
|
||||
case <-subscriber.closed:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) alreadySentLocked(frame *execstream.Frame) bool {
|
||||
return isReplayOutputFrame(frame) &&
|
||||
frame.Watermark != 0 &&
|
||||
frame.Watermark <= subscriber.sentWatermark
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) markSentLocked(frame *execstream.Frame) *execstream.Frame {
|
||||
frame = cloneExecFrame(frame)
|
||||
if isReplayOutputFrame(frame) && frame.Watermark > subscriber.sentWatermark {
|
||||
subscriber.sentWatermark = frame.Watermark
|
||||
}
|
||||
|
||||
return frame
|
||||
}
|
||||
|
||||
func (subscriber *execSessionSubscriber) close() {
|
||||
subscriber.closeOnce.Do(func() {
|
||||
close(subscriber.closed)
|
||||
subscriber.sendMu.Lock()
|
||||
close(subscriber.frames)
|
||||
subscriber.sendMu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
type execSession struct {
|
||||
key execSessionKey
|
||||
command string
|
||||
exec sshExecRunner
|
||||
transport net.Conn
|
||||
registry *execSessionRegistry
|
||||
exitTTL time.Duration
|
||||
policy execSessionPolicy
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
mu sync.Mutex
|
||||
stdin io.WriteCloser
|
||||
stdinClosed bool
|
||||
subscribers map[*execSessionSubscriber]struct{}
|
||||
replay execReplayBuffer
|
||||
started bool
|
||||
finished bool
|
||||
closed bool
|
||||
expiryTimer *time.Timer
|
||||
|
||||
startOnce sync.Once
|
||||
done chan struct{}
|
||||
doneOnce sync.Once
|
||||
}
|
||||
|
||||
func newExecSession(
|
||||
key execSessionKey,
|
||||
command string,
|
||||
exec sshExecRunner,
|
||||
transport net.Conn,
|
||||
registry *execSessionRegistry,
|
||||
exitTTL time.Duration,
|
||||
policy execSessionPolicy,
|
||||
) *execSession {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
return newExecSessionWithContext(
|
||||
ctx,
|
||||
cancel,
|
||||
key,
|
||||
command,
|
||||
exec,
|
||||
transport,
|
||||
registry,
|
||||
exitTTL,
|
||||
policy,
|
||||
)
|
||||
}
|
||||
|
||||
func newExecSessionWithContext(
|
||||
ctx context.Context,
|
||||
cancel context.CancelFunc,
|
||||
key execSessionKey,
|
||||
command string,
|
||||
exec sshExecRunner,
|
||||
transport net.Conn,
|
||||
registry *execSessionRegistry,
|
||||
exitTTL time.Duration,
|
||||
policy execSessionPolicy,
|
||||
) *execSession {
|
||||
if ctx == nil || cancel == nil {
|
||||
ctx, cancel = context.WithCancel(context.Background())
|
||||
}
|
||||
|
||||
session := &execSession{
|
||||
key: key,
|
||||
command: command,
|
||||
exec: exec,
|
||||
transport: transport,
|
||||
registry: registry,
|
||||
exitTTL: exitTTL,
|
||||
policy: policy,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
stdin: exec.Stdin(),
|
||||
subscribers: map[*execSessionSubscriber]struct{}{},
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
return session
|
||||
}
|
||||
|
||||
func (session *execSession) commandMatches(command string) bool {
|
||||
return command == "" || session.command == command
|
||||
}
|
||||
|
||||
func (session *execSession) start() {
|
||||
session.startOnce.Do(func() {
|
||||
session.mu.Lock()
|
||||
if session.closed {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
session.started = true
|
||||
session.mu.Unlock()
|
||||
|
||||
go session.run()
|
||||
})
|
||||
}
|
||||
|
||||
func (session *execSession) closeIfUnused() {
|
||||
session.mu.Lock()
|
||||
unused := !session.started && len(session.subscribers) == 0
|
||||
session.mu.Unlock()
|
||||
|
||||
if unused {
|
||||
session.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) attach() (*execSessionSubscriber, error) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if session.closed {
|
||||
return nil, errors.New("exec session is closed")
|
||||
}
|
||||
|
||||
subscriber := newExecSessionSubscriber()
|
||||
session.subscribers[subscriber] = struct{}{}
|
||||
|
||||
return subscriber, nil
|
||||
}
|
||||
|
||||
func (session *execSession) detach(subscriber *execSessionSubscriber) {
|
||||
if session.policy.closeOnDetach {
|
||||
session.close()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
session.detachLocked(subscriber)
|
||||
}
|
||||
|
||||
func (session *execSession) detachLocked(subscriber *execSessionSubscriber) {
|
||||
if _, ok := session.subscribers[subscriber]; !ok {
|
||||
return
|
||||
}
|
||||
|
||||
delete(session.subscribers, subscriber)
|
||||
subscriber.close()
|
||||
}
|
||||
|
||||
func (session *execSession) writeStdin(data []byte) error {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
if session.stdin == nil || session.stdinClosed {
|
||||
return errors.New("this exec session has no stdin enabled or it is already closed")
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
if err := session.stdin.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
session.stdinClosed = true
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
_, err := session.stdin.Write(data)
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
func (session *execSession) ack(watermark uint64) {
|
||||
if !session.policy.replayEnabled {
|
||||
return
|
||||
}
|
||||
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
session.replay.ack(watermark)
|
||||
}
|
||||
|
||||
func (session *execSession) sendHistory(
|
||||
subscriber *execSessionSubscriber,
|
||||
watermark uint64,
|
||||
) {
|
||||
if !session.policy.replayEnabled {
|
||||
return
|
||||
}
|
||||
|
||||
session.mu.Lock()
|
||||
|
||||
if _, ok := session.subscribers[subscriber]; !ok {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
subscriber.sendMu.Lock()
|
||||
frames := session.replay.replayAfter(watermark, nil)
|
||||
frames = append(frames, &execstream.Frame{
|
||||
Type: execstream.FrameTypeNoMoreHistory,
|
||||
Watermark: session.replay.nextWatermark,
|
||||
})
|
||||
session.mu.Unlock()
|
||||
|
||||
ok := subscriber.sendHistory(frames)
|
||||
subscriber.sendMu.Unlock()
|
||||
|
||||
if !ok {
|
||||
session.dropSubscriber(subscriber)
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) close() {
|
||||
session.mu.Lock()
|
||||
if session.closed {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
session.closed = true
|
||||
if session.expiryTimer != nil {
|
||||
session.expiryTimer.Stop()
|
||||
session.expiryTimer = nil
|
||||
}
|
||||
|
||||
subscribers := session.takeSubscribersLocked()
|
||||
session.mu.Unlock()
|
||||
|
||||
closeSubscribers(subscribers)
|
||||
|
||||
session.cancel()
|
||||
_ = session.exec.Close()
|
||||
if session.transport != nil {
|
||||
_ = session.transport.Close()
|
||||
}
|
||||
if session.registry != nil {
|
||||
session.registry.remove(session.key, session)
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) run() {
|
||||
outgoingFrames := make(chan *execstream.Frame)
|
||||
runErrCh := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
runErrCh <- session.exec.Run(session.ctx, session.command, outgoingFrames)
|
||||
close(outgoingFrames)
|
||||
}()
|
||||
|
||||
for frame := range outgoingFrames {
|
||||
session.recordFrame(frame)
|
||||
}
|
||||
|
||||
runErr := <-runErrCh
|
||||
if runErr != nil && !errors.Is(runErr, context.Canceled) {
|
||||
session.recordFrame(&execstream.Frame{
|
||||
Type: execstream.FrameTypeError,
|
||||
Error: runErr.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
session.markFinished()
|
||||
}
|
||||
|
||||
func (session *execSession) recordFrame(frame *execstream.Frame) {
|
||||
session.mu.Lock()
|
||||
|
||||
if session.closed {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if session.policy.replayEnabled {
|
||||
frame = session.replay.append(frame)
|
||||
} else {
|
||||
frame = cloneExecFrame(frame)
|
||||
}
|
||||
|
||||
subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
|
||||
for subscriber := range session.subscribers {
|
||||
subscribers = append(subscribers, subscriber)
|
||||
}
|
||||
session.mu.Unlock()
|
||||
|
||||
for _, subscriber := range subscribers {
|
||||
if !subscriber.enqueue(frame) {
|
||||
session.dropSubscriber(subscriber)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) markFinished() {
|
||||
session.mu.Lock()
|
||||
if session.finished {
|
||||
session.mu.Unlock()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
session.finished = true
|
||||
shouldClose := !session.policy.retainAfterExit
|
||||
if !session.closed && session.policy.retainAfterExit {
|
||||
session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire)
|
||||
}
|
||||
|
||||
var subscribers []*execSessionSubscriber
|
||||
if shouldClose {
|
||||
subscribers = session.takeSubscribersLocked()
|
||||
}
|
||||
session.mu.Unlock()
|
||||
|
||||
closeSubscribers(subscribers)
|
||||
|
||||
session.doneOnce.Do(func() {
|
||||
close(session.done)
|
||||
})
|
||||
|
||||
if shouldClose {
|
||||
session.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) expire() {
|
||||
session.close()
|
||||
}
|
||||
|
||||
func (session *execSession) takeSubscribersLocked() []*execSessionSubscriber {
|
||||
subscribers := make([]*execSessionSubscriber, 0, len(session.subscribers))
|
||||
for subscriber := range session.subscribers {
|
||||
subscribers = append(subscribers, subscriber)
|
||||
}
|
||||
session.subscribers = map[*execSessionSubscriber]struct{}{}
|
||||
|
||||
return subscribers
|
||||
}
|
||||
|
||||
func closeSubscribers(subscribers []*execSessionSubscriber) {
|
||||
for _, subscriber := range subscribers {
|
||||
subscriber.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (session *execSession) dropSubscriber(subscriber *execSessionSubscriber) {
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
|
||||
session.detachLocked(subscriber)
|
||||
}
|
||||
|
||||
func cloneExecFrame(frame *execstream.Frame) *execstream.Frame {
|
||||
if frame == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
clone := *frame
|
||||
if frame.Data != nil {
|
||||
clone.Data = append([]byte(nil), frame.Data...)
|
||||
}
|
||||
if frame.Exit != nil {
|
||||
exit := *frame.Exit
|
||||
clone.Exit = &exit
|
||||
}
|
||||
|
||||
return &clone
|
||||
}
|
||||
|
||||
func execFrameSize(frame *execstream.Frame) int {
|
||||
if frame == nil {
|
||||
return 0
|
||||
}
|
||||
|
||||
return len(frame.Data) + len(frame.Error) + 16
|
||||
}
|
||||
|
||||
func isReplayOutputFrame(frame *execstream.Frame) bool {
|
||||
if frame == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
switch frame.Type {
|
||||
case execstream.FrameTypeStdout,
|
||||
execstream.FrameTypeStderr,
|
||||
execstream.FrameTypeExit,
|
||||
execstream.FrameTypeError:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user