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) 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) // 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) 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) 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 }