iop/apps/edge/cmd/edge/console_events.go

package main

import (
	"fmt"
	"io"
	"strings"
	"sync"

	"go.uber.org/zap"

	edgenode "iop/apps/edge/internal/node"
	iop "iop/proto/gen/iop"
)

type consoleEventRouter struct {
	mu       sync.Mutex
	out      io.Writer
	registry *edgenode.Registry
	logger   *zap.Logger
	waiters  map[string]chan *iop.RunEvent
	streams  map[string]*consoleResponseStream
}

func newConsoleEventRouter(out io.Writer, registry *edgenode.Registry, logger *zap.Logger) *consoleEventRouter {
	return &consoleEventRouter{
		out:      out,
		registry: registry,
		logger:   logger,
		waiters:  make(map[string]chan *iop.RunEvent),
		streams:  make(map[string]*consoleResponseStream),
	}
}

func (r *consoleEventRouter) Register(runID string) (<-chan *iop.RunEvent, func()) {
	ch := make(chan *iop.RunEvent, 4096)
	r.mu.Lock()
	r.waiters[runID] = ch
	r.mu.Unlock()

	return ch, func() {
		r.mu.Lock()
		if current, ok := r.waiters[runID]; ok && current == ch {
			delete(r.waiters, runID)
			close(ch)
		}
		r.mu.Unlock()
	}
}

func (r *consoleEventRouter) Handle(event *iop.RunEvent) {
	runID := event.GetRunId()

	r.mu.Lock()
	if waiter, ok := r.waiters[runID]; ok {
		select {
		case waiter <- event:
		default:
			if r.logger != nil {
				r.logger.Warn("console run event dropped", zap.String("run_id", runID), zap.String("type", event.GetType()))
			}
		}
		r.mu.Unlock()
		return
	}
	r.printAsyncLocked(event)
	r.mu.Unlock()
}

func (r *consoleEventRouter) nodeLabel(event *iop.RunEvent) string {
	nodeID := event.GetNodeId()
	if nodeID == "" {
		return "unknown"
	}
	if r.registry != nil {
		if entry, ok := r.registry.Get(nodeID); ok && entry.Alias != "" {
			return entry.Alias
		}
	}
	return nodeID
}

func (r *consoleEventRouter) printAsyncLocked(event *iop.RunEvent) {
	if r.out == nil {
		return
	}

	runID := event.GetRunId()
	label := r.nodeLabel(event)
	switch event.GetType() {
	case "start":
		fmt.Fprintf(r.out, "[node-%s-event] start run_id=%s session=%s background=%v\n", label, runID, event.GetSessionId(), event.GetBackground())
	case "delta":
		r.responseStream(runID, label).Write(event.GetDelta())
	case "complete":
		r.finishStream(runID, label)
		fmt.Fprintf(r.out, "[node-%s-event] complete run_id=%s detail=%q\n", label, runID, event.GetMessage())
	case "cancelled":
		r.finishStreamIfStarted(runID)
		fmt.Fprintf(r.out, "[node-%s-event] cancelled run_id=%s\n", label, runID)
	case "error":
		r.finishStreamIfStarted(runID)
		fmt.Fprintf(r.out, "[node-%s-event] error run_id=%s detail=%q\n", label, runID, event.GetError())
	default:
		fmt.Fprintf(r.out, "[node-%s-event] %s run_id=%s detail=%q\n", label, event.GetType(), runID, event.GetMessage())
	}
}

func (r *consoleEventRouter) responseStream(runID, label string) *consoleResponseStream {
	if s, ok := r.streams[runID]; ok {
		return s
	}
	s := newConsoleResponseStream(r.out, fmt.Sprintf("[node-%s-message] ", label))
	r.streams[runID] = s
	return s
}

func (r *consoleEventRouter) finishStream(runID, label string) {
	if s, ok := r.streams[runID]; ok {
		s.Finish()
		delete(r.streams, runID)
		return
	}
	newConsoleResponseStream(r.out, fmt.Sprintf("[node-%s-message] ", label)).Finish()
}

func (r *consoleEventRouter) finishStreamIfStarted(runID string) {
	if s, ok := r.streams[runID]; ok {
		s.FinishIfStarted()
		delete(r.streams, runID)
	}
}

type consoleResponseStream struct {
	out              io.Writer
	prefix           string
	started          bool
	endedWithNewline bool
}

func newConsoleResponseStream(out io.Writer, prefix string) *consoleResponseStream {
	return &consoleResponseStream{out: out, prefix: prefix, endedWithNewline: true}
}

func (s *consoleResponseStream) Write(delta string) {
	if s.out == nil || delta == "" {
		return
	}
	if !s.started {
		fmt.Fprint(s.out, s.prefix)
		s.started = true
	}
	fmt.Fprint(s.out, delta)
	s.endedWithNewline = strings.HasSuffix(delta, "\n")
}

func (s *consoleResponseStream) Finish() {
	if s.out == nil {
		return
	}
	if !s.started {
		fmt.Fprintf(s.out, "%s<empty>\n", s.prefix)
		return
	}
	if !s.endedWithNewline {
		fmt.Fprintln(s.out)
	}
}

func (s *consoleResponseStream) FinishIfStarted() {
	if s.started {
		s.Finish()
	}
}