github.com/kubeshop/testkube@v1.17.23/pkg/triggers/lease_backend.go (about)

     1  package triggers
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  	"go.mongodb.org/mongo-driver/bson"
    10  	"go.mongodb.org/mongo-driver/mongo"
    11  	"go.mongodb.org/mongo-driver/mongo/options"
    12  )
    13  
    14  const (
    15  	mongoCollectionTriggersLease = "triggers"
    16  	documentType                 = "lease"
    17  )
    18  
    19  // LeaseBackend does a check and set operation on the Lease object in the defined data source
    20  //
    21  //go:generate mockgen -destination=./mock_lease_backend.go -package=triggers "github.com/kubeshop/testkube/pkg/triggers" LeaseBackend
    22  type LeaseBackend interface {
    23  	// TryAcquire tries to acquire lease from underlying datastore
    24  	TryAcquire(ctx context.Context, id, clusterID string) (leased bool, err error)
    25  }
    26  
    27  type AcquireAlwaysLeaseBackend struct{}
    28  
    29  func NewAcquireAlwaysLeaseBackend() *AcquireAlwaysLeaseBackend {
    30  	return &AcquireAlwaysLeaseBackend{}
    31  }
    32  
    33  func (b *AcquireAlwaysLeaseBackend) TryAcquire(ctx context.Context, id, clusterID string) (leased bool, err error) {
    34  	return true, nil
    35  }
    36  
    37  type MongoLeaseBackend struct {
    38  	coll *mongo.Collection
    39  }
    40  
    41  func NewMongoLeaseBackend(db *mongo.Database) *MongoLeaseBackend {
    42  	return &MongoLeaseBackend{coll: db.Collection(mongoCollectionTriggersLease)}
    43  }
    44  
    45  func (b *MongoLeaseBackend) TryAcquire(ctx context.Context, id, clusterID string) (leased bool, err error) {
    46  	leaseMongoID := newLeaseMongoID(clusterID)
    47  	currentLease, err := b.findOrInsertCurrentLease(ctx, leaseMongoID, id, clusterID)
    48  	if err != nil {
    49  		return false, err
    50  	}
    51  
    52  	acquired, renewable := leaseStatus(currentLease, id, clusterID)
    53  	switch {
    54  	case acquired:
    55  		return true, nil
    56  	case !renewable:
    57  		return false, nil
    58  	}
    59  
    60  	acquiredAt := currentLease.AcquiredAt
    61  	if currentLease.Identifier != id {
    62  		acquiredAt = time.Now()
    63  	}
    64  	newLease, err := b.tryUpdateLease(ctx, leaseMongoID, id, clusterID, acquiredAt)
    65  	if err != nil {
    66  		return false, err
    67  	}
    68  	acquired, _ = leaseStatus(newLease, id, clusterID)
    69  
    70  	return acquired, nil
    71  }
    72  
    73  func (b *MongoLeaseBackend) findOrInsertCurrentLease(ctx context.Context, leaseMongoID, id, clusterID string) (*Lease, error) {
    74  	res := b.coll.FindOne(ctx, bson.M{"_id": leaseMongoID})
    75  	if res.Err() == mongo.ErrNoDocuments {
    76  		lease, err := b.insertLease(ctx, leaseMongoID, id, clusterID)
    77  		if err != nil {
    78  			return nil, err
    79  		}
    80  		return lease, err
    81  	} else if res.Err() != nil {
    82  		return nil, errors.Wrap(res.Err(), "error finding lease document in mongo")
    83  	}
    84  
    85  	var receivedLease MongoLease
    86  	if err := res.Decode(&receivedLease); err != nil {
    87  		return nil, errors.Wrap(err, "error decoding lease mongo document")
    88  	}
    89  
    90  	return &receivedLease.Lease, nil
    91  }
    92  
    93  func (b *MongoLeaseBackend) insertLease(ctx context.Context, leaseMongoID, id, clusterID string) (*Lease, error) {
    94  	lease := NewLease(id, clusterID)
    95  	_, err := b.coll.InsertOne(ctx, bson.M{"_id": leaseMongoID, "lease": *lease})
    96  	if err != nil {
    97  		return nil, errors.Wrap(err, "error inserting lease document into mongo")
    98  	}
    99  	return lease, nil
   100  }
   101  
   102  func (b *MongoLeaseBackend) tryUpdateLease(ctx context.Context, leaseMongoID, id, clusterID string, acquiredAt time.Time) (*Lease, error) {
   103  	newLease := Lease{
   104  		Identifier: id,
   105  		ClusterID:  clusterID,
   106  		AcquiredAt: acquiredAt,
   107  		RenewedAt:  time.Now(),
   108  	}
   109  	newMongoLease := MongoLease{
   110  		_id:   leaseMongoID,
   111  		Lease: newLease,
   112  	}
   113  
   114  	after := options.After
   115  	opts := options.FindOneAndUpdateOptions{
   116  		ReturnDocument: &after,
   117  	}
   118  	res := b.coll.FindOneAndUpdate(
   119  		ctx,
   120  		bson.M{"_id": leaseMongoID},
   121  		bson.M{"$set": newMongoLease},
   122  		&opts,
   123  	)
   124  	if res.Err() != nil {
   125  		return nil, errors.Wrap(res.Err(), "error finding and updating mongo db document")
   126  	}
   127  	var updatedLease MongoLease
   128  	if err := res.Decode(&updatedLease); err != nil {
   129  		return nil, errors.Wrap(err, "error unmarshalling returned lease mongo document")
   130  	}
   131  
   132  	return &updatedLease.Lease, nil
   133  }
   134  
   135  func leaseStatus(lease *Lease, id, clusterID string) (acquired bool, renewable bool) {
   136  	if lease == nil {
   137  		return false, false
   138  	}
   139  	maxLeaseDurationStaleness := time.Now().Add(-defaultMaxLeaseDuration)
   140  	isLeaseExpired := lease.RenewedAt.Before(maxLeaseDurationStaleness)
   141  	isMyLease := lease.Identifier == id && lease.ClusterID == clusterID
   142  	switch {
   143  	case isLeaseExpired:
   144  		acquired = false
   145  		renewable = true
   146  	case isMyLease:
   147  		acquired = true
   148  		renewable = false
   149  	default:
   150  		acquired = false
   151  		renewable = false
   152  	}
   153  	return
   154  }
   155  
   156  func newLeaseMongoID(clusterID string) string {
   157  	return fmt.Sprintf("%s-%s", documentType, clusterID)
   158  }
   159  
   160  type MongoLease struct {
   161  	_id string `bson:"_id"`
   162  	Lease
   163  }