iop/apps/edge/internal/opsconsole/events.go

package opsconsole

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

	"go.uber.org/zap"

	edgenode "iop/apps/edge/internal/node"
	eventpkg "iop/packages/events"
	iop "iop/proto/gen/iop"
)

// EventRouter renders run and node lifecycle events to the ops console.
// Foreground runs tracked via TrackForeground are suppressed from the
// asynchronous stream so the SendRun loop owns their output.
type EventRouter struct {
	mu         sync.Mutex
	out        io.Writer
	registry   *edgenode.Registry
	logger     *zap.Logger
	foreground map[string]struct{}
	streams    map[string]*ResponseStream
}

func NewEventRouter(out io.Writer, registry *edgenode.Registry, logger *zap.Logger) *EventRouter {
	return &EventRouter{
		out:        out,
		registry:   registry,
		logger:     logger,
		foreground: make(map[string]struct{}),
		streams:    make(map[string]*ResponseStream),
	}
}

func (r *EventRouter) TrackForeground(runID string) func() {
	r.mu.Lock()
	r.foreground[runID] = struct{}{}
	r.mu.Unlock()

	return func() {
		r.mu.Lock()
		delete(r.foreground, runID)
		r.mu.Unlock()
	}
}

func (r *EventRouter) Handle(event *iop.RunEvent) {
	if event == nil || !event.GetBackground() {
		return
	}

	runID := event.GetRunId()

	r.mu.Lock()
	if _, ok := r.foreground[runID]; ok {
		r.mu.Unlock()
		return
	}
	r.printAsyncLocked(event)
	r.mu.Unlock()
}

func (r *EventRouter) HandleNodeEvent(event *iop.EdgeNodeEvent) {
	r.mu.Lock()
	r.printNodeEventLocked(event)
	r.mu.Unlock()
}

func (r *EventRouter) Drain(ctx context.Context, runEvents <-chan *iop.RunEvent, nodeEvents <-chan *iop.EdgeNodeEvent) {
	for {
		select {
		case <-ctx.Done():
			return
		case event, ok := <-runEvents:
			if !ok {
				runEvents = nil
			} else {
				r.Handle(event)
			}
		case event, ok := <-nodeEvents:
			if !ok {
				nodeEvents = nil
			} else {
				r.HandleNodeEvent(event)
			}
		}
		if runEvents == nil && nodeEvents == nil {
			return
		}
	}
}

func (r *EventRouter) NodeLabel(event *iop.RunEvent) string {
	if alias := event.GetNodeAlias(); alias != "" {
		return alias
	}
	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 *EventRouter) edgeNodeLabel(event *iop.EdgeNodeEvent) string {
	if event.GetAlias() != "" {
		return event.GetAlias()
	}
	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 *EventRouter) printNodeEventLocked(event *iop.EdgeNodeEvent) {
	if r.out == nil {
		return
	}

	label := r.edgeNodeLabel(event)
	detail := transportCloseDetail(event.GetMetadata())
	switch event.GetType() {
	case eventpkg.TypeNodeConnected:
		fmt.Fprintf(r.out, "[node-%s-event] connected reason=%q%s\n", label, event.GetReason(), detail)
	case eventpkg.TypeNodeDisconnected:
		fmt.Fprintf(r.out, "[node-%s-event] disconnected reason=%q%s\n", label, event.GetReason(), detail)
	default:
		fmt.Fprintf(r.out, "[node-%s-event] %s reason=%q%s\n", label, event.GetType(), event.GetReason(), detail)
	}
}

func transportCloseDetail(metadata map[string]string) string {
	if metadata == nil {
		return ""
	}
	detail := ""
	if reason := metadata[eventpkg.MetadataTransportCloseReason]; reason != "" {
		detail += fmt.Sprintf(" transport_close_reason=%q", reason)
	}
	if err := metadata[eventpkg.MetadataTransportCloseError]; err != "" {
		detail += fmt.Sprintf(" transport_close_error=%q", err)
	}
	return detail
}

func (r *EventRouter) 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 *EventRouter) responseStream(runID, label string) *ResponseStream {
	if s, ok := r.streams[runID]; ok {
		return s
	}
	s := NewResponseStream(r.out, fmt.Sprintf("[node-%s-message] ", label))
	r.streams[runID] = s
	return s
}

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

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

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

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

func (s *ResponseStream) 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 *ResponseStream) 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 *ResponseStream) FinishIfStarted() {
	if s.started {
		s.Finish()
	}
}