feat: delay lease acquisition by N poll requests (#23)

This is a more simplistic approach to https://github.com/ankorstore/gh-action-mq-lease-service/pull/21

Signed-off-by: Matthias Riegler <[email protected]>

Author: xvzf

internal/config/load.go

@@ -28,9 +28,5 @@ func load(path string, config interface{}) error {
 		return err
 	}
 
-	if err := yaml.Unmarshal([]byte(templated), config); err != nil {
-		return err
-	}
-
-	return nil
+	return yaml.Unmarshal([]byte(templated), config)
 }

internal/config/server/latest/server_schema.go

@@ -13,4 +13,6 @@ type GithubRepositoryConfig struct {
 	StabilizeDuration    int    `yaml:"stabilize_duration_seconds"`
 	TTL                  int    `yaml:"ttl_seconds"`
 	ExpectedRequestCount int    `yaml:"expected_request_count"`
+	// DelayLEaseASsignmentBy is the number of times a lease can be delayed before it is assigned.
+	DelayLeaseAssignmentBy int `yaml:"delay_lease_assignment_by"`
 }

internal/lease/leaseprovider.go

@@ -29,6 +29,7 @@ type ProviderOpts struct {
 	StabilizeDuration    time.Duration
 	TTL                  time.Duration
 	ExpectedRequestCount int
+	DelayAssignmentCount int
 	ID                   string
 	Clock                clock.PassiveClock
 	Storage              storage.Storage[*ProviderState]
@@ -46,11 +47,12 @@ const (
 )
 
 type Request struct {
-	HeadSHA    string  `json:"head_sha"`
-	HeadRef    string  `json:"head_ref"`
-	Priority   int     `json:"priority"`
-	Status     *string `json:"status,omitempty"`
-	lastSeenAt *time.Time
+	HeadSHA          string  `json:"head_sha"`
+	HeadRef          string  `json:"head_ref"`
+	Priority         int     `json:"priority"`
+	Status           *string `json:"status,omitempty"`
+	lastSeenAt       *time.Time
+	acquireCountdown *int
 }
 
 type StackedPullRequest struct {
@@ -75,6 +77,7 @@ func (lr *Request) MarshalZerologObject(e *zerolog.Event) {
 	e.Str("lease_request_head_sha", lr.HeadSHA).
 		Str("lease_request_head_ref", lr.HeadRef).
 		Int("lease_request_priority", lr.Priority).
+		Int("lease_request_acquire_countdown", pointer.IntDeref(lr.acquireCountdown, 0)).
 		Str("lease_request_status", status)
 }
 
@@ -457,12 +460,27 @@ func (lp *leaseProviderImpl) evaluateRequest(ctx context.Context, req *Request)
 
 	// Got the max priority, now check if we are the winner
 	if req.Priority == maxPriority {
-		req.Status = pointer.String(StatusAcquired)
-		lp.state.acquired = req
+
+		// In order to prevent race conditions, there's the option to delay the lock acquisition
+		// This is useful when the lock is acquired by a CI job that is canceled or restarted. There can be a short delay.
+		req.acquireCountdown = pointer.Int(pointer.IntDeref(req.acquireCountdown, lp.opts.DelayAssignmentCount+1) - 1)
+		if *req.acquireCountdown > 0 {
+			log.Ctx(ctx).
+				Debug().
+				EmbedObject(req).
+				Msg("Delaying lock acquisition")
+			return req
+		}
+
 		log.Ctx(ctx).
 			Debug().
 			EmbedObject(req).
 			Msg("Current lease request has the higher priority. It then acquires the lock")
+
+		// Acquire lease
+		req.Status = pointer.String(StatusAcquired)
+		lp.state.acquired = req
+
 		log.Ctx(ctx).
 			Info().
 			EmbedObject(req).

internal/lease/leaseprovider_test.go

@@ -453,7 +453,7 @@ type clearTestFakeStorage struct{ state *ProviderState }
 func (s *clearTestFakeStorage) Init() error                                   { return nil }
 func (s *clearTestFakeStorage) Close() error                                  { return nil }
 func (s *clearTestFakeStorage) Hydrate(context.Context, *ProviderState) error { return nil }
-func (s *clearTestFakeStorage) Save(ctx context.Context, obj *ProviderState) error {
+func (s *clearTestFakeStorage) Save(_ context.Context, obj *ProviderState) error {
 	s.state = obj
 	return nil
 }
@@ -743,3 +743,33 @@ func Test_leaseProviderImpl__FullLoop_ReleaseFromInvalidHeadSHA(t *testing.T) {
 	})
 	assert.Error(t, err)
 }
+
+func Test_leaseProviderImpl__FullLoop_DelayedAcquisition(t *testing.T) {
+	lp := NewLeaseProvider(ProviderOpts{TTL: 1 * time.Hour, StabilizeDuration: time.Minute, ExpectedRequestCount: 2, DelayAssignmentCount: 2})
+
+	req1 := &Request{
+		HeadSHA:  "sha1",
+		Priority: 1,
+	}
+	req2 := &Request{
+		HeadSHA:  "sha2",
+		Priority: 2,
+	}
+
+	// Inject the two requests
+	req1, err := lp.Acquire(context.Background(), req1)
+	assert.NoError(t, err)
+	assert.Equal(t, StatusPending, *req1.Status)
+
+	// this should be delayed and only acquire after the 3rd request
+	req2, err = lp.Acquire(context.Background(), req2)
+	assert.NoError(t, err)
+	assert.Equal(t, StatusPending, *req2.Status)
+	req2, err = lp.Acquire(context.Background(), req2)
+	assert.NoError(t, err)
+	assert.Equal(t, StatusPending, *req2.Status)
+	// Now it should acquire the lease
+	req2, err = lp.Acquire(context.Background(), req2)
+	assert.NoError(t, err)
+	assert.Equal(t, StatusAcquired, *req2.Status)
+}

internal/lease/leaseproviderorchestrator.go

@@ -54,6 +54,7 @@ func NewProviderOrchestrator(opts NewProviderOrchestratorOpts) ProviderOrchestra
 			StabilizeDuration:    time.Second * time.Duration(repository.StabilizeDuration),
 			TTL:                  time.Second * time.Duration(repository.TTL),
 			ExpectedRequestCount: repository.ExpectedRequestCount,
+			DelayAssignmentCount: repository.DelayLeaseAssignmentBy,
 			ID:                   key,
 			Clock:                opts.Clock,
 			Storage:              opts.Storage,

internal/storage/storage.go

@@ -133,7 +133,7 @@ func (s *storageImpl[T]) Hydrate(ctx context.Context, defaultObj T) error {
 
 // Save store the provided object in the storage
 // the provided object should at least be able to return a non-null and unique Identifier (via the GetIdentifier() method)
-func (s *storageImpl[T]) Save(ctx context.Context, obj T) error {
+func (s *storageImpl[T]) Save(_ context.Context, obj T) error {
 	var err error
 	id := obj.GetIdentifier()
 	b, err := obj.Marshal()

pkg/util/logger/logger.go

@@ -18,7 +18,7 @@ type AppInfo interface {
 // locationHook adds the log location
 type locationHook struct{}
 
-func (h locationHook) Run(e *zerolog.Event, l zerolog.Level, msg string) {
+func (h locationHook) Run(e *zerolog.Event, _ zerolog.Level, _ string) {
 	_, file, line, ok := runtime.Caller(3)
 	if ok {
 		e.Str("location", fmt.Sprintf("%s:%d", file, line))