JCS-pub/common/pkgs/ioswitch/exec/executor.go

package exec

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"github.com/hashicorp/go-multierror"
	"gitlink.org.cn/cloudream/common/pkgs/future"
	"gitlink.org.cn/cloudream/common/utils/lo2"
)

type binding struct {
	ID       VarID
	Callback *future.SetValueFuture[VarValue]
}

type freeVar struct {
	ID    VarID
	Value VarValue
}

type Executor struct {
	plan     Plan
	location Location
	vars     map[VarID]freeVar
	bindings []*binding
	lock     sync.Mutex
	store    map[string][]VarValue
}

func NewExecutor(plan Plan, loc Location) *Executor {
	planning := Executor{
		plan:     plan,
		location: loc,
		vars:     make(map[VarID]freeVar),
		store:    make(map[string][]VarValue),
	}

	return &planning
}

func (s *Executor) Plan() *Plan {
	return &s.plan
}

func (s *Executor) Location() Location {
	return s.location
}

func (s *Executor) Run(ctx *ExecContext) (ExecutorResult, error) {
	c, cancel := context.WithCancel(ctx.Context)
	ctx = &ExecContext{
		Context: c,
		Values:  ctx.Values,
	}

	defer cancel()

	err := s.runOps(s.plan.Ops, ctx, cancel)
	if err != nil {
		return ExecutorResult{}, err
	}

	return ExecutorResult{Stored: s.store}, nil
}

func (s *Executor) runOps(ops []Op, ctx *ExecContext, cancel context.CancelFunc) error {
	lock := sync.Mutex{}

	var err error

	var wg sync.WaitGroup
	wg.Add(len(ops))
	for i, arg := range ops {
		go func(op Op, index int) {
			defer wg.Done()

			if e := op.Execute(ctx, s); e != nil {
				lock.Lock()
				// 尽量不记录 Canceled 错误，除非没有其他错误
				if errors.Is(e, context.Canceled) {
					if err == nil {
						err = context.Canceled
					}
				} else {
					err = multierror.Append(err, fmt.Errorf("%T: %w", op, e))
				}
				lock.Unlock()

				cancel()
			}
		}(arg, i)
	}
	wg.Wait()

	return err
}

func (s *Executor) BindVar(ctx context.Context, id VarID) (VarValue, error) {
	s.lock.Lock()

	gv, ok := s.vars[id]
	if ok {
		delete(s.vars, id)
		s.lock.Unlock()
		return gv.Value, nil
	}

	callback := future.NewSetValue[VarValue]()
	s.bindings = append(s.bindings, &binding{
		ID:       id,
		Callback: callback,
	})

	s.lock.Unlock()
	return callback.Wait(ctx)
}

func (s *Executor) PutVar(id VarID, value VarValue) *Executor {
	s.lock.Lock()
	defer s.lock.Unlock()

	for ib, b := range s.bindings {
		if b.ID != id {
			continue
		}

		b.Callback.SetValue(value)
		s.bindings = lo2.RemoveAt(s.bindings, ib)

		return s
	}

	// 如果没有绑定，则直接放入变量表中
	s.vars[id] = freeVar{ID: id, Value: value}
	return s
}

func (s *Executor) Store(key string, val VarValue) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.store[key] = append(s.store[key], val)
}

type ExecutorResult struct {
	Stored map[string][]VarValue
}

func BindVar[T VarValue](e *Executor, ctx context.Context, id VarID) (T, error) {
	v, err := e.BindVar(ctx, id)
	if err != nil {
		var def T
		return def, err
	}

	ret, ok := v.(T)
	if !ok {
		var def T
		return def, fmt.Errorf("binded var %v is %T, not %T", id, v, def)
	}

	return ret, nil
}

func BindArray[T VarValue](e *Executor, ctx context.Context, ids []VarID) ([]T, error) {
	ret := make([]T, len(ids))
	for i := range ids {
		v, err := e.BindVar(ctx, ids[i])
		if err != nil {
			return nil, err
		}

		v2, ok := v.(T)
		if !ok {
			var def T
			return nil, fmt.Errorf("binded var %v is %T, not %T", ids[i], v, def)
		}

		ret[i] = v2
	}
	return ret, nil
}

func PutArray[T VarValue](e *Executor, ids []VarID, values []T) {
	for i := range ids {
		e.PutVar(ids[i], values[i])
	}
}

type Location struct {
	IsDriver   bool
	WorkerName string
}