alt/services/api/internal/workerclient/client.go

package workerclient

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/url"
	"strconv"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"

	altv1 "git.toki-labs.com/toki/alt/packages/contracts/gen/go/alt/v1"
	apiContracts "git.toki-labs.com/toki/alt/services/api/internal/contracts"
	protoSocket "git.toki-labs.com/toki/proto-socket/go"
)

var (
	ErrUnavailable = errors.New("worker is not available")
	ErrTimeout     = errors.New("worker request timeout")
)

type WorkerClient interface {
	Connect(ctx context.Context) error
	Close() error
	IsConnected() bool
	Hello(ctx context.Context, req *altv1.HelloRequest) (*altv1.HelloResponse, error)

	// Backtest command/query surface. The API forwards client requests onto the
	// worker unchanged, so each method mirrors a contract request/response pair.
	StartBacktest(ctx context.Context, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error)
	GetBacktestRun(ctx context.Context, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error)
	ListBacktestRuns(ctx context.Context, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error)
	GetBacktestRunDetail(ctx context.Context, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error)
	GetBacktestResult(ctx context.Context, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error)
	CompareBacktestRuns(ctx context.Context, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error)

	// Paper trading command/query surface. The API forwards client requests onto
	// the worker unchanged; paper execution and state stay worker-owned.
	StartPaperTrading(ctx context.Context, req *altv1.StartPaperTradingRequest) (*altv1.StartPaperTradingResponse, error)
	GetPaperTradingState(ctx context.Context, req *altv1.GetPaperTradingStateRequest) (*altv1.GetPaperTradingStateResponse, error)

	// Market query surface. The API validates request shape and forwards market
	// reads to the worker-owned storage boundary.
	ListInstruments(ctx context.Context, req *altv1.ListInstrumentsRequest) (*altv1.ListInstrumentsResponse, error)
	ListBars(ctx context.Context, req *altv1.ListBarsRequest) (*altv1.ListBarsResponse, error)

	// Market import command surface. Import execution runs inside the worker; the
	// API only validates shape and forwards the command.
	ImportDailyBars(ctx context.Context, req *altv1.ImportDailyBarsRequest) (*altv1.ImportDailyBarsResponse, error)
}

type socketClient struct {
	socketURL string
	mu        sync.RWMutex
	wsClient  *protoSocket.WsClient
}

func New(socketURL string) WorkerClient {
	return &socketClient{
		socketURL: socketURL,
	}
}

func (c *socketClient) IsConnected() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.wsClient != nil && c.wsClient.IsAlive()
}

func (c *socketClient) Connect(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.wsClient != nil && c.wsClient.IsAlive() {
		return nil
	}

	u, err := url.Parse(c.socketURL)
	if err != nil {
		return fmt.Errorf("invalid worker socket URL %q: %w", c.socketURL, err)
	}

	host, portStr, err := net.SplitHostPort(u.Host)
	if err != nil {
		host = u.Host
		if u.Scheme == "wss" {
			portStr = "443"
		} else {
			portStr = "80"
		}
	}

	port, err := strconv.Atoi(portStr)
	if err != nil {
		return fmt.Errorf("invalid worker port %q: %w", portStr, err)
	}

	path := u.Path
	if path == "" {
		path = "/"
	}

	var wsClient *protoSocket.WsClient
	if u.Scheme == "wss" {
		wsClient, err = protoSocket.DialWssWithHeartbeat(ctx, host, port, path, nil, 30, 10, apiContracts.ParserMap())
	} else {
		wsClient, err = protoSocket.DialWsWithHeartbeat(ctx, host, port, path, 30, 10, apiContracts.ParserMap())
	}

	if err != nil {
		return fmt.Errorf("%w: %v", ErrUnavailable, err)
	}

	c.wsClient = wsClient
	return nil
}

func (c *socketClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.wsClient != nil {
		err := c.wsClient.Close()
		c.wsClient = nil
		return err
	}
	return nil
}

func (c *socketClient) Hello(ctx context.Context, req *altv1.HelloRequest) (*altv1.HelloResponse, error) {
	return sendTyped[*altv1.HelloRequest, *altv1.HelloResponse](c, ctx, req)
}

func (c *socketClient) StartBacktest(ctx context.Context, req *altv1.StartBacktestRequest) (*altv1.StartBacktestResponse, error) {
	return sendTyped[*altv1.StartBacktestRequest, *altv1.StartBacktestResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestRun(ctx context.Context, req *altv1.GetBacktestRunRequest) (*altv1.GetBacktestRunResponse, error) {
	return sendTyped[*altv1.GetBacktestRunRequest, *altv1.GetBacktestRunResponse](c, ctx, req)
}

func (c *socketClient) ListBacktestRuns(ctx context.Context, req *altv1.ListBacktestRunsRequest) (*altv1.ListBacktestRunsResponse, error) {
	return sendTyped[*altv1.ListBacktestRunsRequest, *altv1.ListBacktestRunsResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestRunDetail(ctx context.Context, req *altv1.GetBacktestRunDetailRequest) (*altv1.GetBacktestRunDetailResponse, error) {
	return sendTyped[*altv1.GetBacktestRunDetailRequest, *altv1.GetBacktestRunDetailResponse](c, ctx, req)
}

func (c *socketClient) GetBacktestResult(ctx context.Context, req *altv1.GetBacktestResultRequest) (*altv1.GetBacktestResultResponse, error) {
	return sendTyped[*altv1.GetBacktestResultRequest, *altv1.GetBacktestResultResponse](c, ctx, req)
}

func (c *socketClient) CompareBacktestRuns(ctx context.Context, req *altv1.CompareBacktestRunsRequest) (*altv1.CompareBacktestRunsResponse, error) {
	return sendTyped[*altv1.CompareBacktestRunsRequest, *altv1.CompareBacktestRunsResponse](c, ctx, req)
}

func (c *socketClient) StartPaperTrading(ctx context.Context, req *altv1.StartPaperTradingRequest) (*altv1.StartPaperTradingResponse, error) {
	return sendTyped[*altv1.StartPaperTradingRequest, *altv1.StartPaperTradingResponse](c, ctx, req)
}

func (c *socketClient) GetPaperTradingState(ctx context.Context, req *altv1.GetPaperTradingStateRequest) (*altv1.GetPaperTradingStateResponse, error) {
	return sendTyped[*altv1.GetPaperTradingStateRequest, *altv1.GetPaperTradingStateResponse](c, ctx, req)
}

func (c *socketClient) ListInstruments(ctx context.Context, req *altv1.ListInstrumentsRequest) (*altv1.ListInstrumentsResponse, error) {
	return sendTyped[*altv1.ListInstrumentsRequest, *altv1.ListInstrumentsResponse](c, ctx, req)
}

func (c *socketClient) ListBars(ctx context.Context, req *altv1.ListBarsRequest) (*altv1.ListBarsResponse, error) {
	return sendTyped[*altv1.ListBarsRequest, *altv1.ListBarsResponse](c, ctx, req)
}

func (c *socketClient) ImportDailyBars(ctx context.Context, req *altv1.ImportDailyBarsRequest) (*altv1.ImportDailyBarsResponse, error) {
	return sendTyped[*altv1.ImportDailyBarsRequest, *altv1.ImportDailyBarsResponse](c, ctx, req)
}

// sendTyped is the shared request path for every worker call. It centralises
// context-cancellation, deadline-derived timeouts, and the unavailable/timeout
// error mapping so each new request method stays a one-line forwarder and the
// behaviour cannot drift between them.
func sendTyped[Req proto.Message, Res proto.Message](c *socketClient, ctx context.Context, req Req) (Res, error) {
	var zero Res

	if err := ctx.Err(); err != nil {
		return zero, mapContextError(err)
	}

	c.mu.RLock()
	client := c.wsClient
	c.mu.RUnlock()

	if client == nil || !client.IsAlive() {
		return zero, ErrUnavailable
	}

	timeout := 5 * time.Second
	if dl, ok := ctx.Deadline(); ok {
		timeout = time.Until(dl)
		if timeout <= 0 {
			return zero, fmt.Errorf("%w: deadline already exceeded", ErrTimeout)
		}
	}

	type result struct {
		res Res
		err error
	}
	ch := make(chan result, 1)

	go func() {
		res, err := protoSocket.SendRequestTyped[Req, Res](&client.Communicator, req, timeout)
		ch <- result{res: res, err: err}
	}()

	select {
	case <-ctx.Done():
		return zero, mapContextError(ctx.Err())
	case r := <-ch:
		if r.err != nil {
			if errors.Is(r.err, protoSocket.ErrNotConnected) {
				return zero, ErrUnavailable
			}
			return zero, fmt.Errorf("%w: %v", ErrTimeout, r.err)
		}
		return r.res, nil
	}
}

// mapContextError normalises context errors into the worker client's error
// vocabulary: cancellation propagates as-is, deadlines surface as ErrTimeout.
func mapContextError(err error) error {
	if errors.Is(err, context.Canceled) {
		return err
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return fmt.Errorf("%w: %v", ErrTimeout, err)
	}
	return err
}