JCS-pub/common/pkgs/ioswitch2/http_hub_worker.go

package ioswitch2

import (
	"context"
	"fmt"
	"io"
	"strconv"

	"gitlink.org.cn/cloudream/common/utils/io2"
	stgglb "gitlink.org.cn/cloudream/jcs-pub/common/globals"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
	hubapi "gitlink.org.cn/cloudream/jcs-pub/hub/sdk/api"
)

type HttpHubWorker struct {
	Hub cortypes.Hub
}

func (w *HttpHubWorker) NewClient() (exec.WorkerClient, error) {
	addressInfo := w.Hub.Address.(*cortypes.HttpAddressInfo)
	baseUrl := "http://" + addressInfo.ExternalIP + ":" + strconv.Itoa(addressInfo.Port)
	config := hubapi.Config{
		URL: baseUrl,
	}
	pool := hubapi.NewPool(&config)
	cli, err := pool.Acquire()
	defer pool.Release(cli)
	if err != nil {
		return nil, err
	}

	return &HttpHubWorkerClient{hubID: w.Hub.HubID, cli: cli}, nil
}

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

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

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

type HttpHubWorkerClient struct {
	hubID cortypes.HubID
	cli   *hubapi.Client
}

func (c *HttpHubWorkerClient) ExecutePlan(ctx context.Context, plan exec.Plan) (exec.ExecutorResult, error) {
	resp, err := c.cli.ExecuteIOPlan(hubapi.ExecuteIOPlanReq{
		Plan:       plan,
		WorkerName: fmt.Sprintf("%v", c.hubID),
	})
	if err != nil {
		return exec.ExecutorResult{}, err
	}

	return resp.Result, nil
}
func (c *HttpHubWorkerClient) SendStream(ctx context.Context, planID exec.PlanID, id exec.VarID, stream io.ReadCloser) error {
	return c.cli.SendStream(hubapi.SendStreamReq{
		SendStreamInfo: hubapi.SendStreamInfo{
			PlanID: planID,
			VarID:  id,
		},
		Stream: io2.CounterCloser(stream, func(cnt int64, err error) {
			if stgglb.Stats.HubTransfer != nil {
				stgglb.Stats.HubTransfer.RecordOutput(c.hubID, cnt, err == nil || err == io.EOF)
			}
		}),
	})
}
func (c *HttpHubWorkerClient) SendVar(ctx context.Context, planID exec.PlanID, id exec.VarID, value exec.VarValue) error {
	return c.cli.SendVar(hubapi.SendVarReq{
		PlanID:   planID,
		VarID:    id,
		VarValue: value,
	})
}
func (c *HttpHubWorkerClient) GetStream(ctx context.Context, planID exec.PlanID, streamID exec.VarID, signalID exec.VarID, signal exec.VarValue) (io.ReadCloser, error) {
	str, err := c.cli.GetStream(hubapi.GetStreamReq{
		PlanID:   planID,
		VarID:    streamID,
		SignalID: signalID,
		Signal:   signal,
	})
	if err != nil {
		return nil, err
	}

	return io2.CounterCloser(str, func(cnt int64, err error) {
		if stgglb.Stats.HubTransfer != nil {
			stgglb.Stats.HubTransfer.RecordInput(c.hubID, cnt, err == nil || err == io.EOF)
		}
	}), nil
}
func (c *HttpHubWorkerClient) GetVar(ctx context.Context, planID exec.PlanID, varID exec.VarID, signalID exec.VarID, signal exec.VarValue) (exec.VarValue, error) {
	resp, err := c.cli.GetVar(hubapi.GetVarReq{
		PlanID:   planID,
		VarID:    varID,
		SignalID: signalID,
		Signal:   signal,
	})
	if err != nil {
		return nil, err
	}

	return resp.Value, err
}
func (c *HttpHubWorkerClient) Close() error {
	//stgglb.HubRPCPool.Release(c.cli)
	return nil
}