JCS-pub/common/pkgs/ioswitch2/ops2/base_store.go

package ops2

import (
	"fmt"
	"io"
	"time"

	"gitlink.org.cn/cloudream/common/pkgs/future"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/common/utils/io2"
	clitypes "gitlink.org.cn/cloudream/jcs-pub/client/types"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/dag"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch/exec"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/ioswitch2"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/pool"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
)

const (
	BaseReadStatsStoreKey = "BaseReadSpeed"
)

func init() {
	exec.UseOp[*BaseWrite]()
	exec.UseOp[*BaseRead]()
	exec.UseOp[*BaseReadDyn]()
	exec.UseVarValue[*BaseReadStatsValue]()
}

type BaseReadStatsValue struct {
	Size     int64
	Time     time.Duration
	Location exec.Location
}

func (v *BaseReadStatsValue) Clone() exec.VarValue {
	return &BaseReadStatsValue{
		Size:     v.Size,
		Time:     v.Time,
		Location: v.Location,
	}
}

type BaseRead struct {
	Output    exec.VarID
	UserSpace clitypes.UserSpaceDetail
	Path      clitypes.JPath
	Option    types.OpenOption
}

func (o *BaseRead) Execute(ctx *exec.ExecContext, e *exec.Executor) error {
	logger.
		WithField("Output", o.Output).
		WithField("UserSpace", o.UserSpace).
		WithField("Path", o.Path).
		Debug("base read")
	defer logger.Debug("base read end")

	stgPool, err := exec.GetValueByType[*pool.Pool](ctx)
	if err != nil {
		return fmt.Errorf("getting storage pool: %w", err)
	}

	store, err := stgPool.GetBaseStore(&o.UserSpace)
	if err != nil {
		return fmt.Errorf("getting base store of storage %v: %w", o.UserSpace, err)
	}

	stream, err := store.Read(o.Path, o.Option)
	if err != nil {
		return fmt.Errorf("reading object %v: %w", o.Path, err)
	}

	startTime := time.Now()
	counter := io2.CounterCloser(stream, func(cnt int64, err error) {
		if err != nil && err != io.EOF {
			return
		}

		// 要注意这个回调一定要在return之前调用
		e.Store(BaseReadStatsStoreKey, &BaseReadStatsValue{
			Size:     cnt,
			Time:     time.Since(startTime),
			Location: e.Location(),
		})
	})

	fut := future.NewSetVoid()
	output := &exec.StreamValue{
		Stream: io2.AfterReadClosed(counter, func(closer io.ReadCloser) {
			fut.SetVoid()
		}),
	}

	e.PutVar(o.Output, output)
	return fut.Wait(ctx.Context)
}

func (o *BaseRead) String() string {
	return fmt.Sprintf("BaseRead(opt=%v) %v:%v -> %v", o.Option, o.UserSpace, o.Path, o.Output)
}

type BaseReadDyn struct {
	UserSpace clitypes.UserSpaceDetail
	Output    exec.VarID
	FileInfo  exec.VarID
	Option    types.OpenOption
}

func (o *BaseReadDyn) Execute(ctx *exec.ExecContext, e *exec.Executor) error {
	logger.
		WithField("Output", o.Output).
		WithField("UserSpace", o.UserSpace).
		WithField("Path", o.FileInfo).
		Debug("base read dynamic")
	defer logger.Debug("base read dynamic end")

	stgPool, err := exec.GetValueByType[*pool.Pool](ctx)
	if err != nil {
		return fmt.Errorf("getting storage pool: %w", err)
	}

	info, err := exec.BindVar[*FileInfoValue](e, ctx.Context, o.FileInfo)
	if err != nil {
		return err
	}

	store, err := stgPool.GetBaseStore(&o.UserSpace)
	if err != nil {
		return fmt.Errorf("getting base store of storage %v: %w", o.UserSpace, err)
	}

	stream, err := store.Read(info.Path, o.Option)
	if err != nil {
		logger.Warnf("reading file %v: %v", info.Path, err)
		return fmt.Errorf("reading object %v: %w", o.FileInfo, err)
	}

	startTime := time.Now()
	counter := io2.CounterCloser(stream, func(cnt int64, err error) {
		if err != nil && err != io.EOF {
			return
		}

		// 要注意这个回调一定要在return之前调用
		e.Store(BaseReadStatsStoreKey, &BaseReadStatsValue{
			Size:     cnt,
			Time:     time.Since(startTime),
			Location: e.Location(),
		})
	})

	fut := future.NewSetVoid()
	output := &exec.StreamValue{
		Stream: io2.AfterReadClosed(counter, func(closer io.ReadCloser) {
			fut.SetVoid()
		}),
	}

	e.PutVar(o.Output, output)
	return fut.Wait(ctx.Context)
}

func (o *BaseReadDyn) String() string {
	return fmt.Sprintf("BaseReadDyn(opt=%v) %v:%v -> %v", o.Option, o.UserSpace, o.FileInfo, o.Output)
}

type BaseWrite struct {
	Input     exec.VarID
	UserSpace clitypes.UserSpaceDetail
	Path      clitypes.JPath
	FileInfo  exec.VarID
	Option    types.WriteOption
}

func (o *BaseWrite) Execute(ctx *exec.ExecContext, e *exec.Executor) error {
	logger.
		WithField("Input", o.Input).
		Debugf("write file to base store")
	defer logger.Debugf("write file to base store finished")

	stgPool, err := exec.GetValueByType[*pool.Pool](ctx)
	if err != nil {
		return fmt.Errorf("getting storage pool: %w", err)
	}

	store, err := stgPool.GetBaseStore(&o.UserSpace)
	if err != nil {
		return fmt.Errorf("getting base store of storage %v: %w", o.UserSpace, err)
	}

	input, err := exec.BindVar[*exec.StreamValue](e, ctx.Context, o.Input)
	if err != nil {
		return err
	}
	defer input.Stream.Close()

	ret, err := store.Write(o.Path, input.Stream, o.Option)
	if err != nil {
		return err
	}

	e.PutVar(o.FileInfo, &FileInfoValue{
		FileInfo: ret,
	})
	return nil
}

func (o *BaseWrite) String() string {
	return fmt.Sprintf("BaseWrite %v -> %v:%v, %v", o.Input, o.UserSpace, o.Path, o.FileInfo)
}

type BaseReadNode struct {
	dag.NodeBase
	From      ioswitch2.From
	UserSpace clitypes.UserSpaceDetail
	Path      clitypes.JPath
	Option    types.OpenOption
}

func (b *GraphNodeBuilder) NewBaseRead(from ioswitch2.From, userSpace clitypes.UserSpaceDetail, path clitypes.JPath, opt types.OpenOption) *BaseReadNode {
	node := &BaseReadNode{
		From:      from,
		UserSpace: userSpace,
		Path:      path,
		Option:    opt,
	}
	b.AddNode(node)

	node.OutputStreams().Init(node, 1)
	return node
}

func (t *BaseReadNode) GetFrom() ioswitch2.From {
	return t.From
}

func (t *BaseReadNode) Output() dag.StreamOutputSlot {
	return dag.StreamOutputSlot{
		Node:  t,
		Index: 0,
	}
}

func (t *BaseReadNode) GenerateOp() (exec.Op, error) {
	return &BaseRead{
		Output:    t.Output().Var().VarID,
		UserSpace: t.UserSpace,
		Path:      t.Path,
		Option:    t.Option,
	}, nil
}

type BaseReadDynNode struct {
	dag.NodeBase
	From      ioswitch2.From
	UserSpace clitypes.UserSpaceDetail
	Option    types.OpenOption
}

func (b *GraphNodeBuilder) NewBaseReadDyn(from ioswitch2.From, userSpace clitypes.UserSpaceDetail, opt types.OpenOption) *BaseReadDynNode {
	node := &BaseReadDynNode{
		From:      from,
		UserSpace: userSpace,
		Option:    opt,
	}
	b.AddNode(node)

	node.OutputStreams().Init(node, 1)
	node.InputValues().Init(1)
	return node
}

func (t *BaseReadDynNode) GetFrom() ioswitch2.From {
	return t.From
}

func (t *BaseReadDynNode) FileInfoSlot() dag.ValueInputSlot {
	return dag.ValueInputSlot{
		Node:  t,
		Index: 0,
	}
}

func (t *BaseReadDynNode) Output() dag.StreamOutputSlot {
	return dag.StreamOutputSlot{
		Node:  t,
		Index: 0,
	}
}

func (t *BaseReadDynNode) GenerateOp() (exec.Op, error) {
	return &BaseReadDyn{
		UserSpace: t.UserSpace,
		Output:    t.Output().Var().VarID,
		FileInfo:  t.FileInfoSlot().Var().VarID,
		Option:    t.Option,
	}, nil
}

type BaseWriteNode struct {
	dag.NodeBase
	To        ioswitch2.To
	UserSpace clitypes.UserSpaceDetail
	Path      clitypes.JPath
	Option    types.WriteOption
}

func (b *GraphNodeBuilder) NewBaseWrite(to ioswitch2.To, userSpace clitypes.UserSpaceDetail, path clitypes.JPath, opt types.WriteOption) *BaseWriteNode {
	node := &BaseWriteNode{
		To:        to,
		UserSpace: userSpace,
		Path:      path,
		Option:    opt,
	}
	b.AddNode(node)

	node.InputStreams().Init(1)
	node.OutputValues().Init(node, 1)
	return node
}

func (t *BaseWriteNode) GetTo() ioswitch2.To {
	return t.To
}

func (t *BaseWriteNode) Input() dag.StreamInputSlot {
	return dag.StreamInputSlot{
		Node:  t,
		Index: 0,
	}
}

func (t *BaseWriteNode) FileInfoVar() dag.ValueOutputSlot {
	return dag.ValueOutputSlot{
		Node:  t,
		Index: 0,
	}
}

func (t *BaseWriteNode) GenerateOp() (exec.Op, error) {
	return &BaseWrite{
		Input:     t.InputStreams().Get(0).VarID,
		UserSpace: t.UserSpace,
		Path:      t.Path,
		FileInfo:  t.FileInfoVar().Var().VarID,
		Option:    t.Option,
	}, nil
}