Files
muxplex/muxterm/pool.go
T

130 lines
2.8 KiB
Go

package main
import (
"fmt"
"os"
"os/exec"
"sync"
"syscall"
"github.com/creack/pty"
)
// Session represents a PTY-backed tmux attach process.
type Session struct {
PTY *os.File
Cmd *exec.Cmd
Cols uint16
Rows uint16
}
// Pool manages named tmux PTY sessions.
type Pool struct {
mu sync.Mutex
sessions map[string]*Session
}
// NewPool returns an empty Pool.
func NewPool() *Pool {
return &Pool{
sessions: make(map[string]*Session),
}
}
// Attach reuses an existing live session or spawns a new tmux attach process.
// Dead sessions are cleaned up before spawning a replacement.
func (p *Pool) Attach(name string, cols, rows uint16) (*Session, error) {
p.mu.Lock()
defer p.mu.Unlock()
if sess, ok := p.sessions[name]; ok {
// Check if process is still alive via Signal(0).
if sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil {
return sess, nil
}
// Dead session — clean up.
sess.PTY.Close()
delete(p.sessions, name)
}
cmd := exec.Command("tmux", "attach", "-t", name)
winSize := &pty.Winsize{Cols: cols, Rows: rows}
ptmx, err := pty.StartWithSize(cmd, winSize)
if err != nil {
return nil, fmt.Errorf("pty start: %w", err)
}
sess := &Session{
PTY: ptmx,
Cmd: cmd,
Cols: cols,
Rows: rows,
}
p.sessions[name] = sess
// Background goroutine: wait for process exit and clean up.
go func() {
cmd.Wait()
p.mu.Lock()
defer p.mu.Unlock()
// Only delete if the map still holds this exact session
// (Attach may have replaced it already).
if cur, ok := p.sessions[name]; ok && cur == sess {
sess.PTY.Close()
delete(p.sessions, name)
}
}()
return sess, nil
}
// Resize changes the PTY window size for a named session.
func (p *Pool) Resize(name string, cols, rows uint16) error {
p.mu.Lock()
defer p.mu.Unlock()
sess, ok := p.sessions[name]
if !ok {
return fmt.Errorf("session %q not found", name)
}
if err := pty.Setsize(sess.PTY, &pty.Winsize{Cols: cols, Rows: rows}); err != nil {
return fmt.Errorf("setsize: %w", err)
}
sess.Cols = cols
sess.Rows = rows
return nil
}
// Get returns the session for name, or nil if not present.
func (p *Pool) Get(name string) *Session {
p.mu.Lock()
defer p.mu.Unlock()
return p.sessions[name]
}
// IsAlive returns true if the named session exists and its process is alive.
func (p *Pool) IsAlive(name string) bool {
p.mu.Lock()
defer p.mu.Unlock()
sess, ok := p.sessions[name]
if !ok {
return false
}
return sess.Cmd.Process != nil && sess.Cmd.Process.Signal(syscall.Signal(0)) == nil
}
// CloseAll sends SIGTERM to all sessions and closes their PTYs.
func (p *Pool) CloseAll() {
p.mu.Lock()
defer p.mu.Unlock()
for name, sess := range p.sessions {
if sess.Cmd.Process != nil {
sess.Cmd.Process.Signal(syscall.SIGTERM)
}
sess.PTY.Close()
delete(p.sessions, name)
}
}