iotlcm/internal/batch/runtime.go

package batch

import (
	"context"
	"strings"
	"time"

	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"

	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model/entity"
)

type Repository interface {
	FindMain(ctx context.Context, batchNo string) (*entity.BatchMain, error)
	MarkProcessing(ctx context.Context, batchNo string, batchType int) (bool, error)
	MarkFailed(ctx context.Context, batchNo string, batchType int, reason string) (bool, error)
	MarkSuccess(ctx context.Context, batchNo string, batchType int, success, fail int, reason string) (bool, error)
}

type Runtime struct {
	repository          Repository
	dispatcher          *Dispatcher
	executionTimeout    time.Duration
	finalizationTimeout time.Duration
}

func NewRuntime(repository Repository, dispatcher *Dispatcher, options ...RuntimeOption) *Runtime {
	runtime := &Runtime{
		repository:          repository,
		dispatcher:          dispatcher,
		executionTimeout:    0,
		finalizationTimeout: 10 * time.Second,
	}
	for _, option := range options {
		option(runtime)
	}
	return runtime
}

type RuntimeOption func(*Runtime)

func WithExecutionTimeout(timeout time.Duration) RuntimeOption {
	return func(runtime *Runtime) {
		if timeout >= 0 {
			runtime.executionTimeout = timeout
		}
	}
}

func WithFinalizationTimeout(timeout time.Duration) RuntimeOption {
	return func(runtime *Runtime) {
		if timeout > 0 {
			runtime.finalizationTimeout = timeout
		}
	}
}

func (r *Runtime) Handle(ctx context.Context, message Message) error {
	if strings.TrimSpace(message.BatchNo) == "" {
		g.Log().Warning(ctx, "batch message dropped, empty batchNo")
		return nil
	}
	executor, ok := r.dispatcher.Executor(message.BatchType)
	if !ok {
		g.Log().Warningf(ctx, "batch message dropped, unsupported batchType:%d, batchNo:%s", message.BatchType, message.BatchNo)
		return nil
	}

	main, err := r.repository.FindMain(ctx, message.BatchNo)
	if err != nil {
		return gerror.Wrap(err, "load batch main failed")
	}
	if main == nil {
		g.Log().Warningf(ctx, "batch message ignored, batch not found, batchNo:%s", message.BatchNo)
		return nil
	}
	if main.BatchType != message.BatchType {
		g.Log().Warningf(ctx, "batch message ignored, type mismatch, batchNo:%s, messageType:%d, storedType:%d", message.BatchNo, message.BatchType, main.BatchType)
		return nil
	}
	if strings.TrimSpace(message.FailureReason) != "" {
		return r.markExecutionFailed(ctx, message, executor, message.FailureReason)
	}
	if main.BatchAsyncStatus != AsyncStatusPending {
		if main.BatchAsyncStatus == AsyncStatusProcessing {
			if message.Retryable {
				return r.reclaimAndExecute(ctx, message, executor, main)
			}
			if processingExecutor, ok := executor.(ProcessingExecutor); ok {
				canExecute, checkErr := processingExecutor.CanExecuteProcessing(ctx, main)
				if checkErr != nil {
					return gerror.Wrap(checkErr, "check processing batch re-entry failed")
				}
				if canExecute {
					if selfClaimer, selfClaims := executor.(SelfClaimingProcessingExecutor); selfClaims && selfClaimer.ClaimsProcessingReentry() {
						return r.execute(ctx, message, executor)
					}
					return r.reclaimAndExecute(ctx, message, executor, main)
				}
			}
		}
		logUnclaimable(ctx, main)
		return nil
	}

	claimed, err := r.repository.MarkProcessing(ctx, message.BatchNo, message.BatchType)
	if err != nil {
		return gerror.Wrap(err, "mark batch processing failed")
	}
	if !claimed {
		current, findErr := r.repository.FindMain(ctx, message.BatchNo)
		if findErr != nil {
			return gerror.Wrap(findErr, "reload batch main after CAS failed")
		}
		if current == nil {
			g.Log().Warningf(ctx, "batch CAS lost, batch not found, batchNo:%s", message.BatchNo)
			return nil
		}
		if current.BatchType != message.BatchType {
			g.Log().Warningf(ctx, "batch CAS lost, type mismatch, batchNo:%s, messageType:%d, storedType:%d", message.BatchNo, message.BatchType, current.BatchType)
			return nil
		}
		logUnclaimable(ctx, current)
		return nil
	}

	g.Log().Infof(ctx, "batch claimed, batchNo:%s, batchType:%d", message.BatchNo, message.BatchType)
	return r.execute(ctx, message, executor)
}

func (r *Runtime) reclaimAndExecute(ctx context.Context, message Message, executor Executor, main *entity.BatchMain) error {
	if claimer, claimSupported := r.repository.(ProcessingClaimer); claimSupported {
		reclaimed, claimErr := claimer.ReclaimProcessing(ctx, message.BatchNo, message.BatchType, main.StartTime)
		if claimErr != nil {
			return gerror.Wrap(claimErr, "claim stale processing batch failed")
		}
		if !reclaimed {
			g.Log().Infof(ctx, "batch re-entry claim lost, batchNo:%s", message.BatchNo)
			return nil
		}
	}
	return r.execute(ctx, message, executor)
}

func (r *Runtime) execute(ctx context.Context, message Message, executor Executor) error {
	executionCtx := ctx
	var cancel context.CancelFunc = func() {}
	if r.executionTimeout > 0 {
		executionCtx, cancel = context.WithTimeout(ctx, r.executionTimeout)
	}
	defer cancel()

	if err := executor.Execute(executionCtx, message.BatchNo); err != nil {
		if IsRetryable(err) {
			return err
		}
		reason := err.Error()
		if r.executionTimeout > 0 && executionCtx.Err() == context.DeadlineExceeded {
			reason = gerror.Newf("batch execution timed out after %s", r.executionTimeout).Error()
		}
		return r.markExecutionFailed(ctx, message, executor, reason)
	}
	if r.executionTimeout > 0 && executionCtx.Err() == context.DeadlineExceeded {
		return r.markExecutionFailed(ctx, message, executor, gerror.Newf("batch execution timed out after %s", r.executionTimeout).Error())
	}
	return nil
}

func (r *Runtime) markExecutionFailed(ctx context.Context, message Message, executor Executor, reason string) error {
	finalizationCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), r.finalizationTimeout)
	defer cancel()

	failed, markErr := r.repository.MarkFailed(finalizationCtx, message.BatchNo, message.BatchType, reason)
	if markErr != nil {
		return &FailureMarkError{Reason: reason, Err: markErr}
	}
	if handler, ok := executor.(ExecutionFailureHandler); ok {
		if err := handler.HandleExecutionFailure(finalizationCtx, message.BatchNo, message.BatchType, reason); err != nil {
			return &FailureMarkError{Reason: reason, Err: err}
		}
	}
	if !failed {
		g.Log().Warningf(ctx, "batch execution failed but status was already finalized, batchNo:%s, reason:%s", message.BatchNo, reason)
		return nil
	}
	g.Log().Errorf(ctx, "batch execution failed and was marked failed, batchNo:%s, reason:%s", message.BatchNo, reason)
	return nil
}

func logUnclaimable(ctx context.Context, main *entity.BatchMain) {
	if main.BatchAsyncStatus == AsyncStatusProcessing {
		g.Log().Infof(ctx, "batch message ignored, batch already processing, batchNo:%s", main.BatchNo)
		return
	}
	g.Log().Infof(ctx, "batch message ignored, batch is terminal or not pending, batchNo:%s, status:%d", main.BatchNo, main.BatchAsyncStatus)
}