JCS-pub/common/pkgs/ioswitchlrc/hub_worker.go

package ioswitchlrc

import (
	"context"
	"io"

	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"
	hubrpc "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc/hub"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
)

// var _ = serder.UseTypeUnionExternallyTagged(types.Ref(types.NewTypeUnion[exec.WorkerInfo](
// 	(*HubWorker)(nil),
// )))

type HubWorker struct {
	Hub     cortypes.Hub
	Address cortypes.GRPCAddressInfo
}

func (w *HubWorker) NewClient() (exec.WorkerClient, error) {
	cli := stgglb.HubRPCPool.Get(stgglb.SelectGRPCAddress(&w.Hub, &w.Address))
	return &HubWorkerClient{cli: cli}, nil
}

func (w *HubWorker) String() string {
	return w.Hub.String()
}

func (w *HubWorker) Equals(worker exec.WorkerInfo) bool {
	aw, ok := worker.(*HubWorker)
	if !ok {
		return false
	}

	return w.Hub.HubID == aw.Hub.HubID
}

type HubWorkerClient struct {
	cli *hubrpc.Client
}

func (c *HubWorkerClient) ExecutePlan(ctx context.Context, plan exec.Plan) (exec.ExecutorResult, error) {
	resp, err := c.cli.ExecuteIOPlan(ctx, &hubrpc.ExecuteIOPlan{Plan: plan})
	if err != nil {
		return exec.ExecutorResult{}, err.ToError()
	}

	return resp.Result, nil
}
func (c *HubWorkerClient) SendStream(ctx context.Context, planID exec.PlanID, id exec.VarID, stream io.ReadCloser) error {
	_, err := c.cli.SendIOStream(ctx, &hubrpc.SendIOStream{
		PlanID: planID,
		VarID:  id,
		Stream: stream,
	})
	return err.ToError()
}
func (c *HubWorkerClient) SendVar(ctx context.Context, planID exec.PlanID, id exec.VarID, value exec.VarValue) error {
	_, err := c.cli.SendIOVar(ctx, &hubrpc.SendIOVar{
		PlanID: planID, VarID: id, Value: value,
	})
	return err.ToError()
}
func (c *HubWorkerClient) GetStream(ctx context.Context, planID exec.PlanID, streamID exec.VarID, signalID exec.VarID, signal exec.VarValue) (io.ReadCloser, error) {
	resp, err := c.cli.GetIOStream(ctx, &hubrpc.GetIOStream{
		PlanID: planID, VarID: streamID, SignalID: signalID, Signal: signal,
	})
	if err != nil {
		return nil, err.ToError()
	}

	return resp.Stream, nil
}
func (c *HubWorkerClient) GetVar(ctx context.Context, planID exec.PlanID, varID exec.VarID, signalID exec.VarID, signal exec.VarValue) (exec.VarValue, error) {
	resp, err := c.cli.GetIOVar(ctx, &hubrpc.GetIOVar{
		PlanID: planID, VarID: varID, SignalID: signalID, Signal: signal,
	})
	if err != nil {
		return nil, err.ToError()
	}
	return resp.Value, nil
}
func (c *HubWorkerClient) Close() error {
	c.cli.Release()
	return nil
}