A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
6.9 kB · 255 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256package controller
import ( "context" "errors" "fmt" "sync" "time"
"connectrpc.com/connect" agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" "forgejo.alexma.top/alexma233/composia/internal/platform/rpcutil" "github.com/google/uuid")
const ( containerLogKindReady = "ready" containerLogKindStart = "start" containerLogKindClose = "close" containerLogKindChunk = "chunk" containerLogKindError = "error" containerLogKindClosed = "closed")
type containerLogTunnelManager struct { mu sync.Mutex tunnels map[string]*agentContainerLogTunnel sessions map[string]*containerLogSession}
type agentContainerLogTunnel struct { nodeID string sendCh chan *agentv1.OpenContainerLogTunnelResponse}
type containerLogSession struct { id string nodeID string incoming chan *agentv1.OpenContainerLogTunnelRequest created time.Time
mu sync.Mutex closed bool}
func newContainerLogTunnelManager() *containerLogTunnelManager { return &containerLogTunnelManager{ tunnels: make(map[string]*agentContainerLogTunnel), sessions: make(map[string]*containerLogSession), }}
func (manager *containerLogTunnelManager) registerTunnel(nodeID string) *agentContainerLogTunnel { tunnel := &agentContainerLogTunnel{ nodeID: nodeID, sendCh: make(chan *agentv1.OpenContainerLogTunnelResponse, 256), } manager.mu.Lock() previous := manager.tunnels[nodeID] manager.tunnels[nodeID] = tunnel if previous != nil { close(previous.sendCh) } manager.mu.Unlock() return tunnel}
func (manager *containerLogTunnelManager) unregisterTunnel(nodeID string, tunnel *agentContainerLogTunnel) { manager.mu.Lock() defer manager.mu.Unlock() if manager.tunnels[nodeID] != tunnel { return } delete(manager.tunnels, nodeID) close(tunnel.sendCh) for sessionID, session := range manager.sessions { if session.nodeID != nodeID { continue } delete(manager.sessions, sessionID) session.close() }}
func (manager *containerLogTunnelManager) hasTunnel(nodeID string) bool { manager.mu.Lock() defer manager.mu.Unlock() _, ok := manager.tunnels[nodeID] return ok}
func (manager *containerLogTunnelManager) openSession(nodeID, containerID, tail string, timestamps bool) (*containerLogSession, error) { manager.mu.Lock() defer manager.mu.Unlock() tunnel := manager.tunnels[nodeID] if tunnel == nil { return nil, fmt.Errorf("node %q has no active container log tunnel", nodeID) } session := &containerLogSession{ id: uuid.NewString(), nodeID: nodeID, incoming: make(chan *agentv1.OpenContainerLogTunnelRequest, 256), created: time.Now().UTC(), } manager.sessions[session.id] = session select { case tunnel.sendCh <- &agentv1.OpenContainerLogTunnelResponse{ SessionId: session.id, Kind: containerLogKindStart, ContainerId: containerID, Tail: tail, Timestamps: timestamps, }: default: delete(manager.sessions, session.id) session.close() return nil, fmt.Errorf("node %q container log tunnel is not accepting messages", nodeID) } return session, nil}
func (manager *containerLogTunnelManager) closeSession(sessionID string) { manager.mu.Lock() defer manager.mu.Unlock() session := manager.sessions[sessionID] if session == nil { return } delete(manager.sessions, sessionID) session.close() if tunnel := manager.tunnels[session.nodeID]; tunnel != nil { select { case tunnel.sendCh <- &agentv1.OpenContainerLogTunnelResponse{SessionId: sessionID, Kind: containerLogKindClose}: default: } }}
func (manager *containerLogTunnelManager) deliverFromAgent(message *agentv1.OpenContainerLogTunnelRequest) { if message != nil && message.GetKind() == containerLogKindReady { return } if message == nil || message.GetSessionId() == "" { return } manager.mu.Lock() session := manager.sessions[message.GetSessionId()] if message.GetKind() == containerLogKindError || message.GetKind() == containerLogKindClosed { delete(manager.sessions, message.GetSessionId()) } manager.mu.Unlock() if session == nil { return } if !session.deliver(message) { manager.closeSession(message.GetSessionId()) return } if message.GetKind() == containerLogKindError || message.GetKind() == containerLogKindClosed { session.close() }}
func (session *containerLogSession) deliver(message *agentv1.OpenContainerLogTunnelRequest) bool { session.mu.Lock() defer session.mu.Unlock() if session.closed { return false } select { case session.incoming <- message: return true default: return false }}
func (session *containerLogSession) close() { session.mu.Lock() defer session.mu.Unlock() if session.closed { return } session.closed = true close(session.incoming)}
func (server *agentReportServer) OpenContainerLogTunnel(ctx context.Context, stream *connect.BidiStream[agentv1.OpenContainerLogTunnelRequest, agentv1.OpenContainerLogTunnelResponse]) error { nodeID, ok := rpcutil.BearerSubject(ctx) if !ok || nodeID == "" { return connect.NewError(connect.CodeUnauthenticated, errors.New("missing bearer subject")) } manager := server.logManager if manager == nil { return connect.NewError(connect.CodeFailedPrecondition, errors.New("container log manager is not configured")) } tunnel := manager.registerTunnel(nodeID) defer manager.unregisterTunnel(nodeID, tunnel)
recvErrCh := make(chan error, 1) go func() { for { message, recvErr := stream.Receive() if recvErr != nil { recvErrCh <- recvErr return } manager.deliverFromAgent(message) } }()
for { select { case <-ctx.Done(): return ctx.Err() case recvErr := <-recvErrCh: if errors.Is(recvErr, context.Canceled) { return nil } return recvErr case message, ok := <-tunnel.sendCh: if !ok { return nil } if err := stream.Send(message); err != nil { return err } } }}
func (server *dockerCommandServer) openContainerLogSession(nodeID, containerID, tail string, timestamps bool) (*containerLogSession, error) { if server.logManager == nil { return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("container log manager is not configured")) } if !server.logManager.hasTunnel(nodeID) { return nil, connect.NewError(connect.CodeFailedPrecondition, fmt.Errorf("node %q has no active container log tunnel", nodeID)) } session, err := server.logManager.openSession(nodeID, containerID, tail, timestamps) if err != nil { return nil, connect.NewError(connect.CodeFailedPrecondition, err) } return session, nil}
func containerLogStreamError(message *agentv1.OpenContainerLogTunnelRequest) error { if message == nil { return connect.NewError(connect.CodeInternal, errors.New("container log stream failed")) } text := message.GetErrorMessage() if text == "" { text = "container log stream failed" } return connect.NewError(dockerQueryConnectCode(message.GetErrorCode()), errors.New(text))}