op-packages/smart-srun/core/internal/control/dispatch.go

package control

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"time"

	"github.com/matthewlu070111/smart-srun/core/internal/domain"
)

// Handler runs one method. Params are the raw request parameters; the handler
// decodes them into its own type with DecodeParams.
type Handler func(ctx context.Context, params json.RawMessage) (any, error)

// Registry binds declared method names to handlers.
//
// Only names from the catalogue can be bound, and each exactly once. The
// dispatcher looks a name up in this map -- it never derives a function, a
// command or a path from the request.
type Registry struct {
	handlers map[string]Handler
}

func NewRegistry() *Registry {
	return &Registry{handlers: make(map[string]Handler, len(catalogue))}
}

// Register binds a handler. It panics on an undeclared or duplicate name:
// both are wiring mistakes that must fail at startup, not serve a request.
func (r *Registry) Register(name string, handler Handler) {
	if _, declared := Lookup(name); !declared {
		panic("control: registering undeclared method " + name)
	}
	if _, duplicate := r.handlers[name]; duplicate {
		panic("control: method registered twice: " + name)
	}
	if handler == nil {
		panic("control: nil handler for " + name)
	}
	r.handlers[name] = handler
}

// Registered reports whether a handler is bound.
func (r *Registry) Registered(name string) bool {
	_, ok := r.handlers[name]
	return ok
}

// Missing lists declared methods with no handler, so a startup check can refuse
// to serve a half-wired protocol instead of returning NotFound at runtime for
// something the client was told exists.
func (r *Registry) Missing() []string {
	var out []string
	for _, name := range MethodNames() {
		if !r.Registered(name) {
			out = append(out, name)
		}
	}
	return out
}

// Dispatch runs one request and always produces a response.
//
// Errors become error responses rather than propagating: the peer is waiting
// for exactly one line, and dropping the connection on an unexpected error
// would leave it to time out with no idea what happened.
func (r *Registry) Dispatch(ctx context.Context, request Request) Response {
	method, declared := Lookup(request.Method)
	if !declared {
		return NewFailure(request.RequestID, domain.Errorf(domain.CodeNotFound,
			"未知方法 %q", request.Method))
	}
	handler, bound := r.handlers[method.Name]
	if !bound {
		return NewFailure(request.RequestID,
			domain.Errorf(domain.CodeUnsupportedCapability,
				"方法 %q 在本次构建中尚未实现", method.Name))
	}

	// Reads are bounded here rather than in each handler, so a read added later
	// cannot forget to be. Mutations and tasks keep the caller's deadline: a
	// login or an install is bounded by its own use case, and cutting one off
	// after a poll's three seconds would look like a network failure.
	if method.Kind == KindRead {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, ReadResponseDeadline)
		defer cancel()
	}

	result, err := handler(ctx, request.Params)
	if err != nil {
		return NewFailure(request.RequestID, err)
	}
	response, err := NewSuccess(request.RequestID, result)
	if err != nil {
		return NewFailure(request.RequestID, err)
	}
	return response
}

// DecodeParams decodes request parameters into a typed value.
//
// Unknown fields are refused. That is what stops a caller from attaching a
// parameter one method honours to another that would ignore it -- the shape of
// what a method accepts is part of its contract, not a suggestion.
func DecodeParams(raw json.RawMessage, target any) error {
	if len(raw) == 0 || string(raw) == "null" {
		return nil
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return domain.Errorf(domain.CodeInvalidArgument,
			"参数无效：%v", err).Wrap(err)
	}
	if _, err := decoder.Token(); !errors.Is(err, io.EOF) {
		return domain.Errorf(domain.CodeInvalidArgument, "params 之后还有内容")
	}
	return nil
}

// ServeConn handles one connection: one request in, one response out.
//
// One request per connection keeps the protocol free of pipelining, ordering
// and half-closed-stream questions that nothing here needs. A peer that
// disconnects before completing its request gets no response, because there is
// nobody left to read one.
func ServeConn(ctx context.Context, registry *Registry, stream io.ReadWriter) error {
	timed, _ := stream.(deadlineStream)
	if timed != nil {
		if err := timed.SetReadDeadline(time.Now().Add(RequestReadDeadline)); err != nil {
			return err
		}
	}

	reader := bufio.NewReaderSize(stream, 4096)
	request, err := ReadRequest(reader, MaxRequestBytes)
	if err != nil {
		if errors.Is(err, ErrClosed) {
			return nil
		}
		// The request id is unknown when parsing failed, so the reply carries
		// an empty one. The peer still learns why it was refused.
		return WriteResponse(stream, NewFailure("", err))
	}

	// The response budget starts once the request is understood and covers the
	// handler and the write together, so a read cannot spend its three seconds
	// working and then take three more to answer.
	respondBy := time.Now().Add(ReadResponseDeadline)
	response := registry.Dispatch(ctx, request)

	if timed != nil {
		if method, declared := Lookup(request.Method); declared && method.Kind == KindRead {
			if err := timed.SetWriteDeadline(respondBy); err != nil {
				return err
			}
		}
	}
	return WriteResponse(stream, response)
}