go-omni/ws/rpc/peer.go

package rpc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/ws"
)

type Handler func(context.Context, json.RawMessage) (any, error)

type transport interface {
	Context() context.Context
	Send(context.Context, ws.Message) error
	TrySend(ws.Message) error
	Close(int, string) error
	Abort()
}

// Peer belongs to exactly one physical session. Construct a new Peer on every
// reconnect. Register routes before passing Handlers to a client/server.
type Peer struct {
	mu       sync.Mutex
	cfg      Config
	conn     transport
	ctx      context.Context
	cancel   context.CancelFunc
	routes   map[string]Handler
	pending  map[string]chan envelope
	incoming map[string]context.CancelFunc
	seq      uint64
	active   int
	closed   bool
	done     chan struct{}
	ready    chan struct{}
	doneOnce sync.Once
}

func New(cfg Config) (*Peer, error) {
	cfg, err := normalize(cfg)
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithCancel(context.Background())
	return &Peer{cfg: cfg, ctx: ctx, cancel: cancel, routes: make(map[string]Handler), pending: make(map[string]chan envelope), incoming: make(map[string]context.CancelFunc), done: make(chan struct{}), ready: make(chan struct{})}, nil
}

func (p *Peer) Handle(method string, handler Handler) error {
	if !validName(method, 128) || handler == nil {
		return errors.New("rpc: valid method and handler required")
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return ErrClosed
	}
	if p.conn != nil {
		return errors.New("rpc: register methods before binding")
	}
	if _, ok := p.routes[method]; ok {
		return errors.New("rpc: duplicate method")
	}
	p.routes[method] = handler
	return nil
}

// Handlers integrates protocol work into the Session completion barrier.
// OnClose waits for all accepted handlers, including handlers ignoring cancel.
// Owners must use a bounded Session.Wait/Shutdown and retain dependencies if it
// times out. Do not call Wait or Shutdown from a registered handler itself.
func (p *Peer) Handlers() ws.Handlers {
	return ws.Handlers{
		OnOpen: func(_ context.Context, s *ws.Session) error {
			if s.Subprotocol() != Subprotocol {
				return ErrProtocol
			}
			return p.bind(s)
		},
		OnMessage: func(_ context.Context, _ *ws.Session, m ws.Message) error { return p.receive(m) },
		OnClose:   func(_ *ws.Session, _ error) { p.Stop(); _ = p.Wait(context.Background()) },
	}
}

func (p *Peer) bind(c transport) error {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return ErrClosed
	}
	if p.conn != nil {
		p.mu.Unlock()
		return errors.New("rpc: peer already bound")
	}
	p.cancel()
	p.ctx, p.cancel = context.WithCancel(c.Context())
	p.conn = c
	close(p.ready)
	ctx := p.ctx
	p.mu.Unlock()
	context.AfterFunc(ctx, p.Stop)
	return nil
}

// Stop cancels pending calls and incoming handlers without waiting for them.
// It does not own the underlying Session; Shutdown additionally closes it.
func (p *Peer) Stop() {
	p.mu.Lock()
	if !p.closed {
		p.closed = true
		p.cancel()
		for _, cancel := range p.incoming {
			cancel()
		}
	}
	if p.active == 0 {
		p.doneOnce.Do(func() { close(p.done) })
	}
	p.mu.Unlock()
}
func (p *Peer) Done() <-chan struct{} { return p.done }

// WaitReady waits for OnOpen to bind the transport. Call and Restore should run
// outside OnOpen: replies are delivered only after that callback returns.
func (p *Peer) WaitReady(ctx context.Context) error {
	select {
	case <-p.ready:
		p.mu.Lock()
		closed := p.closed
		p.mu.Unlock()
		if closed {
			return ErrClosed
		}
		return nil
	case <-p.done:
		return ErrClosed
	case <-ctx.Done():
		return ctx.Err()
	}
}
func (p *Peer) Wait(ctx context.Context) error {
	select {
	case <-p.done:
		return nil
	default:
	}
	select {
	case <-p.done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}
func (p *Peer) Shutdown(ctx context.Context) error {
	p.Stop()
	p.mu.Lock()
	c := p.conn
	p.mu.Unlock()
	if c != nil {
		_ = c.Close(1000, "")
	}
	return p.Wait(ctx)
}

// Call waits for the peer's reply, bounded by ctx and Config.CallTimeout.
// Cancellation is best effort on the wire; a timeout never proves rollback.
func (p *Peer) Call(ctx context.Context, method string, input, out any) error {
	if !validName(method, 128) {
		return errors.New("rpc: invalid method")
	}
	params, err := json.Marshal(input)
	if err != nil {
		return fmt.Errorf("rpc: encode input: %w", err)
	}
	ctx, cancel := context.WithTimeout(ctx, p.cfg.CallTimeout)
	defer cancel()
	if err := ctx.Err(); err != nil {
		return err
	}
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return ErrClosed
	}
	if p.conn == nil {
		p.mu.Unlock()
		return ErrNotConnected
	}
	if len(p.pending) >= p.cfg.MaxPending {
		p.mu.Unlock()
		return ErrBusy
	}
	p.seq++
	id := strconv.FormatUint(p.seq, 36)
	ch := make(chan envelope, 1)
	p.pending[id] = ch
	c, work := p.conn, p.ctx
	p.mu.Unlock()
	defer func() { p.mu.Lock(); delete(p.pending, id); p.mu.Unlock() }()
	deadline, _ := ctx.Deadline()
	budget := max(int64(1), time.Until(deadline).Milliseconds())
	if err = p.send(ctx, c, envelope{Version: 1, Type: "call", ID: id, Method: method, Params: params, TimeoutMS: budget}); err != nil {
		return err
	}
	select {
	case reply := <-ch:
		if reply.Error != nil {
			return reply.Error
		}
		if out == nil {
			return nil
		}
		return json.Unmarshal(reply.Result, out)
	case <-ctx.Done():
		// Never block a caller waiting for room to send cancellation. The remote
		// handler is independently bounded by the transmitted deadline.
		_ = p.try(c, envelope{Version: 1, Type: "cancel", ID: id})
		return ctx.Err()
	case <-work.Done():
		return ErrClosed
	}
}

// Notify confirms a local write, not remote processing. Notifications are
// handled with the same concurrency and timeout limits as incoming calls.
func (p *Peer) Notify(ctx context.Context, method string, input any) error {
	if !validName(method, 128) {
		return errors.New("rpc: invalid method")
	}
	params, err := json.Marshal(input)
	if err != nil {
		return err
	}
	p.mu.Lock()
	c, closed := p.conn, p.closed
	p.mu.Unlock()
	if closed {
		return ErrClosed
	}
	if c == nil {
		return ErrNotConnected
	}
	ctx, cancel := context.WithTimeout(ctx, p.cfg.WriteTimeout)
	defer cancel()
	return p.send(ctx, c, envelope{Version: 1, Type: "event", Method: method, Params: params})
}

func (p *Peer) encode(e envelope) (ws.Message, error) {
	data, err := json.Marshal(e)
	if err != nil {
		return ws.Message{}, err
	}
	if len(data) > p.cfg.MaxMessageBytes {
		return ws.Message{}, ws.ErrMessageTooLarge
	}
	return ws.Message{Type: ws.Text, Data: data}, nil
}
func (p *Peer) send(ctx context.Context, c transport, e envelope) error {
	m, err := p.encode(e)
	if err != nil {
		return err
	}
	return c.Send(ctx, m)
}
func (p *Peer) try(c transport, e envelope) error {
	m, err := p.encode(e)
	if err != nil {
		return err
	}
	return c.TrySend(m)
}

func (p *Peer) receive(m ws.Message) error {
	if m.Type != ws.Text {
		return ErrProtocol
	}
	e, err := decode(m.Data, p.cfg.MaxMessageBytes)
	if err != nil {
		return err
	}
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return ErrClosed
	}
	if p.conn == nil {
		p.mu.Unlock()
		return ErrNotConnected
	}
	switch e.Type {
	case "reply":
		if ch := p.pending[e.ID]; ch != nil {
			select {
			case ch <- e:
			default:
			}
		}
		p.mu.Unlock()
		return nil
	case "cancel":
		if cancel := p.incoming[e.ID]; cancel != nil {
			cancel()
		}
		p.mu.Unlock()
		return nil
	}
	c := p.conn
	if e.Type == "call" {
		if _, exists := p.incoming[e.ID]; exists {
			p.mu.Unlock()
			return ErrProtocol
		}
	}
	handler := p.routes[e.Method]
	if handler == nil {
		p.mu.Unlock()
		if e.Type == "event" {
			return ErrProtocol
		}
		return p.try(c, envelope{Version: 1, Type: "reply", ID: e.ID, Error: &Error{"METHOD_NOT_FOUND", "method not found"}})
	}
	if p.active >= p.cfg.MaxConcurrent {
		p.mu.Unlock()
		if e.Type == "event" {
			return ErrBusy
		}
		return p.try(c, envelope{Version: 1, Type: "reply", ID: e.ID, Error: &Error{"BUSY", "peer capacity reached"}})
	}
	timeout := p.cfg.CallTimeout
	if e.TimeoutMS > 0 {
		timeout = min(timeout, time.Duration(e.TimeoutMS)*time.Millisecond)
	}
	ctx, cancel := context.WithTimeout(p.ctx, timeout)
	if e.Type == "call" {
		p.incoming[e.ID] = cancel
	}
	p.active++
	p.mu.Unlock()
	go p.execute(ctx, cancel, c, e, handler)
	return nil
}

func invoke(ctx context.Context, h Handler, params json.RawMessage) (value any, err error) {
	defer func() {
		if recover() != nil {
			value = nil
			err = errors.New("rpc: handler panic")
		}
	}()
	return h(ctx, params)
}
func publicError(err error) (result *Error) {
	result = &Error{"INTERNAL", "request failed"}
	defer func() {
		if recover() != nil {
			result = &Error{"INTERNAL", "request failed"}
		}
	}()
	if errors.Is(err, context.DeadlineExceeded) {
		return &Error{"DEADLINE_EXCEEDED", "request deadline exceeded"}
	}
	if errors.Is(err, context.Canceled) {
		return &Error{"CANCELED", "request canceled"}
	}
	var exposed *Error
	if errors.As(err, &exposed) && validError(exposed) {
		return &Error{exposed.Code, exposed.Message}
	}
	return &Error{"INTERNAL", "request failed"}
}

func marshalResult(value any) (raw json.RawMessage, err error) {
	defer func() {
		if recover() != nil {
			raw = nil
			err = errors.New("rpc: result encoding failed")
		}
	}()
	return json.Marshal(value)
}
func (p *Peer) execute(ctx context.Context, cancel context.CancelFunc, c transport, e envelope, h Handler) {
	defer func() {
		cancel()
		p.mu.Lock()
		delete(p.incoming, e.ID)
		p.active--
		if p.closed && p.active == 0 {
			p.doneOnce.Do(func() { close(p.done) })
		}
		p.mu.Unlock()
	}()
	value, err := invoke(ctx, h, e.Params)
	if e.Type == "event" {
		if err != nil && p.cfg.OnEventError != nil {
			func() { defer func() { _ = recover() }(); p.cfg.OnEventError(ctx, err) }()
		}
		return
	}
	if err == nil {
		err = ctx.Err()
	}
	reply := envelope{Version: 1, Type: "reply", ID: e.ID}
	if err == nil {
		reply.Result, err = marshalResult(value)
	}
	if err != nil {
		reply.Result = nil
		reply.Error = publicError(err)
	}
	writeCtx, writeCancel := context.WithTimeout(p.ctx, p.cfg.WriteTimeout)
	defer writeCancel()
	if err = p.send(writeCtx, c, reply); errors.Is(err, ws.ErrMessageTooLarge) {
		reply.Result = nil
		reply.Error = &Error{"RESULT_TOO_LARGE", "result exceeds message limit"}
		err = p.send(writeCtx, c, reply)
	}
	if err != nil {
		c.Abort()
	}
}