Maven/internal/worker/server.go

// worker/server.go — the worker-side listener + accept loop.
//
// Each call is dispatched to a single handler (Transcriber OR Synthesizer,
// depending on what the module process);
// the other verb returns ErrUnknownMethod — a stt process won't serve
// synthesize. robust to a misconfigured client (the daemon wiring chooses
// which module to dial; mixing the two is a config error caught cleanly by
// the wire, not a runtime goroutine panic). One Server per module process.
//
// Socket perms mirror ipc.Server: dir 0700, socket 0600 ⇒ same unix user.
// The module has no key, so the floor is "same user"; the wg/mTLS layers
// are out of scope here (this socket never crosses the network radius —
// it's local-only, point-to-point between two processes on the box).
package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"sync"
	"sync/atomic"

	"golang.org/x/sys/unix"
)

// Server — a worker module process's listener. Wires either a Transcriber,
// a Synthesizer, or both (the both case is unusual; the daemon typically
// runs two separate module processes). The unset verb returns
// ErrUnknownMethod.
type Server struct {
	t Transcriber
	s Synthesizer

	path string
	ln   net.Listener

	wg   sync.WaitGroup
	done chan struct{}

	// connCount — assigned per accepted conn, used in logs to distinguish
	// concurrent connections. Monotonic; not load-bearing for correctness.
	connCount atomic.Uint64
}

// NewServer builds a Server with a Transcriber. The caller wires a
// Synthesizer via SetSynthesizer if this process serves tts. Use
// NewSynthesizerServer for the tts-only case (mirrors this constructor).
func NewServer(path string, t Transcriber) *Server {
	return &Server{t: t, path: path, done: make(chan struct{})}
}

// NewSynthesizerServer builds a Server with a Synthesizer (the tts module).
func NewSynthesizerServer(path string, s Synthesizer) *Server {
	return &Server{s: s, path: path, done: make(chan struct{})}
}

// SetSynthesizer wires the synthesize verb on a Transcriber-built Server.
// Used only when one process serves both (non-default; the daemon prefers
// two separate processes per the restart-free / fail-independent invariant).
func (srv *Server) SetSynthesizer(s Synthesizer) { srv.s = s }

// Listen binds the unix socket with 0700 dir + 0600 socket perms (same floor
// as internal/ipc). A stale socket at path is removed first so the worker
// process restarts cleanly after a crash, no manual cleanup needed.
func (srv *Server) Listen() error {
	_ = os.Remove(srv.path)
	if err := os.MkdirAll(parentDir(srv.path), 0o700); err != nil {
		return fmt.Errorf("worker: mkdir socket dir: %w", err)
	}
	oldMask := unix.Umask(0o077)
	ln, err := net.Listen("unix", srv.path)
	unix.Umask(oldMask)
	if err != nil {
		return fmt.Errorf("worker: listen %s: %w", srv.path, err)
	}
	if err := os.Chmod(srv.path, 0o600); err != nil {
		_ = ln.Close()
		_ = os.Remove(srv.path)
		return fmt.Errorf("worker: chmod socket: %w", err)
	}
	srv.ln = ln
	return nil
}

// Path returns the bound socket path (after Listen; "" before).
func (srv *Server) Path() string { return srv.path }

// Serve accepts connections until the listener closes. Each connection is
// served in its own goroutine; a panicking handler tears down only that conn
// (the rest of the module keeps serving, restart-free per spec).
func (srv *Server) Serve() error {
	if srv.ln == nil {
		return fmt.Errorf("worker: serve before listen")
	}
	for {
		c, err := srv.ln.Accept()
		if err != nil {
			select {
			case <-srv.done:
				return nil
			default:
				return fmt.Errorf("worker: accept: %w", err)
			}
		}
		srv.wg.Add(1)
		go func(c net.Conn) {
			defer srv.wg.Done()
			defer c.Close()
			srv.serveConn(c)
		}(c)
	}
}

func (srv *Server) serveConn(c net.Conn) {
	id := srv.connCount.Add(1)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	for {
		var req Request
		if err := readFrame(c, &req); err != nil {
			return // EOF / malformed ⇒ end this conn
		}
		result, err := srv.safeDispatch(ctx, c.RemoteAddr(), id, req)
		resp := Response{}
		if err != nil {
			resp.Error = rpcErr(err)
		} else {
			resp.Result = result
		}
		if err := writeFrame(c, resp); err != nil {
			return
		}
	}
}

func (srv *Server) safeDispatch(ctx context.Context, addr net.Addr, id uint64, req Request) (result json.RawMessage, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("worker: panic dispatching %s (conn %d): %v", req.Method, id, r)
		}
	}()
	return srv.dispatch(ctx, req)
}

func (srv *Server) dispatch(ctx context.Context, req Request) (json.RawMessage, error) {
	switch req.Method {
	case MethodTranscribe:
		if srv.t == nil {
			return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
		}
		var p TranscribeReq
		if err := unmarshalParams(req.Params, &p); err != nil {
			return nil, err
		}
		resp, err := srv.t.Transcribe(ctx, p)
		if err != nil {
			return nil, err
		}
		return marshalResult(resp), nil

	case MethodSynthesize:
		if srv.s == nil {
			return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
		}
		var p SynthesizeReq
		if err := unmarshalParams(req.Params, &p); err != nil {
			return nil, err
		}
		resp, err := srv.s.Synthesize(ctx, p)
		if err != nil {
			return nil, err
		}
		return marshalResult(resp), nil

	default:
		return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
	}
}

// Close stops accepting and waits for in-flight connections to drain. The
// socket file is removed so a restart can rebind cleanly. Idempotent.
func (srv *Server) Close() error {
	select {
	case <-srv.done:
		return nil
	default:
		close(srv.done)
	}
	if srv.ln == nil {
		return nil
	}
	err := srv.ln.Close()
	srv.wg.Wait()
	_ = os.Remove(srv.path)
	return err
}

func unmarshalParams(raw json.RawMessage, v any) error {
	if len(raw) == 0 {
		raw = []byte("null")
	}
	if err := json.Unmarshal(raw, v); err != nil {
		return fmt.Errorf("%w: %v", ErrBadParams, err)
	}
	return nil
}

func marshalResult(v any) json.RawMessage {
	if v == nil {
		return json.RawMessage("null")
	}
	b, _ := json.Marshal(v)
	return b
}

func parentDir(p string) string {
	for i := len(p) - 1; i >= 0; i-- {
		if p[i] == '/' {
			if i == 0 {
				return "/"
			}
			return p[:i]
		}
	}
	return "."
}