github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/utils/transaction.go (about)

     1  package utils
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/projecteru2/core/log"
     8  )
     9  
    10  // ContextFunc .
    11  type contextFunc = func(context.Context) error
    12  
    13  // Txn provides unified API to perform txn
    14  func Txn(ctx context.Context, cond contextFunc, then contextFunc, rollback func(context.Context, bool) error, ttl time.Duration) (txnErr error) {
    15  	var condErr, thenErr error
    16  	txnCtx, txnCancel := context.WithTimeout(ctx, ttl)
    17  	defer txnCancel()
    18  	logger := log.WithFunc("utils.Txn")
    19  	defer func() { // rollback
    20  		txnErr = condErr
    21  		if txnErr == nil {
    22  			txnErr = thenErr
    23  		}
    24  		if txnErr == nil {
    25  			return
    26  		}
    27  		if rollback == nil {
    28  			logger.Warn(ctx, "txn failed but no rollback function")
    29  			return
    30  		}
    31  
    32  		logger.Error(ctx, txnErr, "txn failed, rolling back")
    33  
    34  		// forbid interrupting rollback
    35  		rollbackCtx, rollBackCancel := context.WithTimeout(NewInheritCtx(ctx), ttl)
    36  		defer rollBackCancel()
    37  		failureByCond := condErr != nil
    38  		if err := rollback(rollbackCtx, failureByCond); err != nil {
    39  			logger.Warnf(ctx, "txn failed but rollback also failed: %+v", err)
    40  		}
    41  	}()
    42  
    43  	// let caller decide process then or not
    44  	if condErr = cond(txnCtx); condErr == nil && then != nil {
    45  		// no rollback and forbid interrupting further process
    46  		thenCtx := txnCtx
    47  		var thenCancel context.CancelFunc
    48  		if rollback == nil {
    49  			thenCtx, thenCancel = context.WithTimeout(NewInheritCtx(ctx), ttl)
    50  			defer thenCancel()
    51  		}
    52  		thenErr = then(thenCtx)
    53  	}
    54  
    55  	return txnErr
    56  }
    57  
    58  // PCR Prepare, Commit, Rollback.
    59  // `prepare` should be a pure calculation process without side effects.
    60  // `commit` writes the calculation result of `prepare` into database.
    61  // if `commit` returns error, `rollback` will be performed.
    62  func PCR(ctx context.Context, prepare func(ctx context.Context) error, commit func(ctx context.Context) error, rollback func(ctx context.Context) error, ttl time.Duration) error {
    63  	return Txn(ctx, prepare, commit, func(ctx context.Context, failureByCond bool) error {
    64  		if !failureByCond {
    65  			return rollback(ctx)
    66  		}
    67  		return nil
    68  	}, ttl)
    69  }