Captioneer/src/app/worker_protocol.go

package app

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"tea.chunkbyte.com/kato/captioneer/src/captions"
	"tea.chunkbyte.com/kato/captioneer/src/config"
)

const (
	modelWorkerEnvironment   = "CAPTIONEER_INTERNAL_MODEL_WORKER"
	modelSettingsEnvironment = "CAPTIONEER_INTERNAL_MODEL_SETTINGS"
	modelProtocolFD          = 3
)

const (
	messageReady       = "ready"
	messageBusy        = "busy"
	messageIdle        = "idle"
	messageCaption     = "caption"
	messageFailure     = "failure"
	messageHeartbeat   = "heartbeat"
	messageSource      = "source"
	messageSpeech      = "speech"
	messageRecognition = "recognition"
	messageOverloaded  = "overloaded"
)

type workerMessage struct {
	Type   string          `json:"type"`
	Event  *captions.Event `json:"event,omitempty"`
	Error  string          `json:"error,omitempty"`
	Value  string          `json:"value,omitempty"`
	Active bool            `json:"active,omitempty"`
}

type workerReporter struct {
	mu      sync.Mutex
	encoder *json.Encoder
}

func (r *workerReporter) send(message workerMessage) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.encoder.Encode(message)
}

type protocolSink struct {
	reporter *workerReporter
}

func (s protocolSink) Publish(event captions.Event) error {
	return s.reporter.send(workerMessage{Type: messageCaption, Event: &event})
}

func (protocolSink) Close() error { return nil }

// IsModelWorker reports whether this process was created as Captioneer's
// private inference worker. The environment marker is intentionally internal.
func IsModelWorker() bool {
	return os.Getenv(modelWorkerEnvironment) == "1"
}

func encodeWorkerSettings(settings config.Settings) (string, error) {
	data, err := json.Marshal(settings)
	if err != nil {
		return "", err
	}
	return base64.RawStdEncoding.EncodeToString(data), nil
}

func decodeWorkerSettings(value string) (config.Settings, error) {
	data, err := base64.RawStdEncoding.DecodeString(value)
	if err != nil {
		return config.Settings{}, fmt.Errorf("decode model worker settings: %w", err)
	}
	var settings config.Settings
	if err := json.Unmarshal(data, &settings); err != nil {
		return config.Settings{}, fmt.Errorf("parse model worker settings: %w", err)
	}
	if err := settings.Validate(); err != nil {
		return config.Settings{}, fmt.Errorf("validate model worker settings: %w", err)
	}
	return settings, nil
}

// RunModelWorker owns all Sherpa/VAD resources in the supervised child. It
// returns a process exit code so command entrypoints can call os.Exit.
func RunModelWorker(diagnostics io.Writer) int {
	protocol := os.NewFile(uintptr(modelProtocolFD), "captioneer-model-protocol")
	if protocol == nil {
		fmt.Fprintln(diagnostics, "captioneer: model worker protocol is unavailable")
		return 1
	}
	defer protocol.Close()
	syscall.CloseOnExec(modelProtocolFD)

	settings, err := decodeWorkerSettings(os.Getenv(modelSettingsEnvironment))
	if err != nil {
		fmt.Fprintf(diagnostics, "captioneer: %v\n", err)
		return 1
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	reporter := &workerReporter{encoder: json.NewEncoder(protocol)}
	hooks := directHooks{
		ready:       func() error { return reporter.send(workerMessage{Type: messageReady}) },
		busy:        func() error { return reporter.send(workerMessage{Type: messageBusy}) },
		idle:        func() error { return reporter.send(workerMessage{Type: messageIdle}) },
		source:      func(value string) error { return reporter.send(workerMessage{Type: messageSource, Value: value}) },
		speech:      func(active bool) error { return reporter.send(workerMessage{Type: messageSpeech, Active: active}) },
		recognition: func(active bool) error { return reporter.send(workerMessage{Type: messageRecognition, Active: active}) },
		overloaded:  func() error { return reporter.send(workerMessage{Type: messageOverloaded}) },
	}

	heartbeatInterval := time.Second
	if settings.ModelTimeout > 0 && settings.ModelTimeout < 4*heartbeatInterval {
		heartbeatInterval = settings.ModelTimeout / 4
		if heartbeatInterval <= 0 {
			heartbeatInterval = time.Millisecond
		}
	}
	heartbeatContext, stopHeartbeat := context.WithCancel(ctx)
	heartbeatDone := make(chan struct{})
	go func() {
		defer close(heartbeatDone)
		ticker := time.NewTicker(heartbeatInterval)
		defer ticker.Stop()
		for {
			select {
			case <-heartbeatContext.Done():
				return
			case <-ticker.C:
				if reporter.send(workerMessage{Type: messageHeartbeat}) != nil {
					return
				}
			}
		}
	}()

	err = runDirect(ctx, settings, protocolSink{reporter: reporter}, diagnostics, hooks)
	stopHeartbeat()
	<-heartbeatDone
	if err == nil || errors.Is(err, context.Canceled) {
		return 0
	}
	_ = reporter.send(workerMessage{Type: messageFailure, Error: err.Error()})
	fmt.Fprintf(diagnostics, "captioneer: model worker: %v\n", err)
	return 1
}