113 lines
2.6 KiB
Go
113 lines
2.6 KiB
Go
package toolsrv
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// ProcessKeeper maintains a running ollie-remote process and respawns it
|
|
// on failure. Multiple agents can Dial() the same keeper to get independent
|
|
// connections. If the process dies, the next Dial() triggers a respawn.
|
|
type ProcessKeeper struct {
|
|
mu sync.Mutex
|
|
ctx context.Context
|
|
proc *Process
|
|
spawn func(context.Context) (*Process, error)
|
|
}
|
|
|
|
// NewProcessKeeper creates a keeper with an initial process and a spawn factory.
|
|
func NewProcessKeeper(ctx context.Context, proc *Process, spawn func(context.Context) (*Process, error)) *ProcessKeeper {
|
|
return &ProcessKeeper{
|
|
ctx: ctx,
|
|
proc: proc,
|
|
spawn: spawn,
|
|
}
|
|
}
|
|
|
|
// Dial returns a new Conn to the managed process. If the process is dead
|
|
// (socket unreachable), it respawns before dialing.
|
|
func (pk *ProcessKeeper) Dial() (*Conn, error) {
|
|
pk.mu.Lock()
|
|
sock := pk.proc.Socket
|
|
pk.mu.Unlock()
|
|
|
|
// Try dialing the existing socket.
|
|
conn, err := Dial(pk.ctx, sock)
|
|
if err == nil {
|
|
return conn, nil
|
|
}
|
|
|
|
// Dial failed — process likely dead. Respawn.
|
|
return pk.respawnAndDial()
|
|
}
|
|
|
|
// Socket returns the current process socket path.
|
|
func (pk *ProcessKeeper) Socket() string {
|
|
pk.mu.Lock()
|
|
defer pk.mu.Unlock()
|
|
return pk.proc.Socket
|
|
}
|
|
|
|
// Process returns the current underlying Process.
|
|
func (pk *ProcessKeeper) Process() *Process {
|
|
pk.mu.Lock()
|
|
defer pk.mu.Unlock()
|
|
return pk.proc
|
|
}
|
|
|
|
// Close shuts down the managed process.
|
|
func (pk *ProcessKeeper) Close() {
|
|
pk.mu.Lock()
|
|
defer pk.mu.Unlock()
|
|
if pk.proc != nil {
|
|
pk.proc.Close()
|
|
}
|
|
}
|
|
|
|
// SetContext updates the keeper's context. Used on session resume to
|
|
// provide a fresh context after the previous one was cancelled on pause.
|
|
func (pk *ProcessKeeper) SetContext(ctx context.Context) {
|
|
pk.mu.Lock()
|
|
defer pk.mu.Unlock()
|
|
pk.ctx = ctx
|
|
}
|
|
|
|
func (pk *ProcessKeeper) respawnAndDial() (*Conn, error) {
|
|
pk.mu.Lock()
|
|
defer pk.mu.Unlock()
|
|
|
|
// Another goroutine may have respawned already — try current socket first.
|
|
if conn, err := Dial(pk.ctx, pk.proc.Socket); err == nil {
|
|
return conn, nil
|
|
}
|
|
|
|
// Clean up old process.
|
|
pk.proc.Close()
|
|
|
|
var lastErr error
|
|
for attempt := 0; attempt < 3; attempt++ {
|
|
if pk.ctx.Err() != nil {
|
|
return nil, pk.ctx.Err()
|
|
}
|
|
proc, err := pk.spawn(pk.ctx)
|
|
if err != nil {
|
|
lastErr = err
|
|
time.Sleep(time.Duration(attempt+1) * 2 * time.Second)
|
|
continue
|
|
}
|
|
pk.proc = proc
|
|
|
|
conn, err := Dial(pk.ctx, proc.Socket)
|
|
if err != nil {
|
|
lastErr = err
|
|
proc.Close()
|
|
time.Sleep(time.Duration(attempt+1) * 2 * time.Second)
|
|
continue
|
|
}
|
|
return conn, nil
|
|
}
|
|
return nil, fmt.Errorf("respawn failed after 3 attempts: %w", lastErr)
|
|
}
|