generated from TBD54566975/tbd-project-template
-
Notifications
You must be signed in to change notification settings - Fork 8
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add leases and switch scheduledtask to use them (#1351)
Leases provide limited exclusive access to a resource for a period of time with automatic renewal. Note: Currently each lease renewal occurs in a dedicated goroutine, via an UPDATE issued every two seconds. An easy future optimisation to reduce the number of updates is to aggregate all UPDATES into a single statement run at the same frequency. See [design doc](https://hackmd.io/@ftl/Sym_GKEb0).
- Loading branch information
1 parent
f8ce99d
commit 4edee31
Showing
19 changed files
with
612 additions
and
116 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| package dal | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "time" | ||
|
|
||
| "github.com/google/uuid" | ||
|
|
||
| "github.com/TBD54566975/ftl/backend/controller/leases" | ||
| "github.com/TBD54566975/ftl/backend/controller/sql" | ||
| "github.com/TBD54566975/ftl/internal/log" | ||
| ) | ||
|
|
||
| const leaseRenewalInterval = time.Second * 2 | ||
|
|
||
| var _ leases.Leaser = (*DAL)(nil) | ||
|
|
||
| // Lease represents a lease that is held by a controller. | ||
| type Lease struct { | ||
| idempotencyKey uuid.UUID | ||
| context any | ||
| key leases.Key | ||
| db *sql.DB | ||
| ttl time.Duration | ||
| errch chan error | ||
| release chan bool | ||
| leak bool // For testing. | ||
| } | ||
|
|
||
| func (l *Lease) String() string { | ||
| return fmt.Sprintf("%s:%s", l.key, l.idempotencyKey) | ||
| } | ||
|
|
||
| // Periodically renew the lease until it is released. | ||
| func (l *Lease) renew(ctx context.Context) { | ||
| defer close(l.errch) | ||
| logger := log.FromContext(ctx).Scope("lease(" + l.key.String() + ")") | ||
| logger.Debugf("Acquired lease %s", l.key) | ||
| for { | ||
| select { | ||
| case <-time.After(leaseRenewalInterval): | ||
| logger.Tracef("Renewing lease %s", l.key) | ||
| ctx, cancel := context.WithTimeout(ctx, leaseRenewalInterval) | ||
| _, err := l.db.RenewLease(ctx, l.ttl, l.idempotencyKey, l.key) | ||
| cancel() | ||
|
|
||
| if err != nil { | ||
| logger.Errorf(err, "Failed to renew lease %s", l.key) | ||
| l.errch <- translatePGError(err) | ||
| return | ||
| } | ||
|
|
||
| case <-l.release: | ||
| if l.leak { // For testing. | ||
| return | ||
| } | ||
| logger.Debugf("Releasing lease %s", l.key) | ||
| _, err := l.db.ReleaseLease(ctx, l.idempotencyKey, l.key) | ||
| l.errch <- translatePGError(err) | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (l *Lease) Release() error { | ||
| close(l.release) | ||
| return <-l.errch | ||
| } | ||
|
|
||
| func (d *DAL) AcquireLease(ctx context.Context, key leases.Key, ttl time.Duration) (leases.Lease, error) { | ||
| if ttl < time.Second*5 { | ||
| return nil, fmt.Errorf("lease TTL must be at least 5 seconds") | ||
| } | ||
| idempotencyKey, err := d.db.NewLease(ctx, key, time.Now().Add(ttl)) | ||
| if err != nil { | ||
| return nil, translatePGError(err) | ||
| } | ||
| lease := &Lease{ | ||
| idempotencyKey: idempotencyKey, | ||
| context: nil, | ||
| key: key, | ||
| db: d.db, | ||
| ttl: ttl, | ||
| release: make(chan bool), | ||
| errch: make(chan error, 1), | ||
| } | ||
| go lease.renew(ctx) | ||
| return lease, nil | ||
| } | ||
|
|
||
| // ExpireLeases expires (deletes) all leases that have expired. | ||
| func (d *DAL) ExpireLeases(ctx context.Context) error { | ||
| count, err := d.db.ExpireLeases(ctx) | ||
| // TODO: Return and log the actual lease keys? | ||
| if count > 0 { | ||
| log.FromContext(ctx).Warnf("Expired %d leases", count) | ||
| } | ||
| return translatePGError(err) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,94 @@ | ||
| package dal | ||
|
|
||
| import ( | ||
| "context" | ||
| "errors" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/alecthomas/assert/v2" | ||
| "github.com/google/uuid" | ||
|
|
||
| "github.com/TBD54566975/ftl/backend/controller/leases" | ||
| "github.com/TBD54566975/ftl/backend/controller/sql" | ||
| "github.com/TBD54566975/ftl/backend/controller/sql/sqltest" | ||
| "github.com/TBD54566975/ftl/internal/log" | ||
| ) | ||
|
|
||
| func leaseExists(t *testing.T, conn sql.DBI, idempotencyKey uuid.UUID, key leases.Key) bool { | ||
| t.Helper() | ||
| var count int | ||
| err := translatePGError(conn. | ||
| QueryRow(context.Background(), "SELECT COUNT(*) FROM leases WHERE idempotency_key = $1 AND key = $2", idempotencyKey, key). | ||
| Scan(&count)) | ||
| if errors.Is(err, ErrNotFound) { | ||
| return false | ||
| } | ||
| assert.NoError(t, err) | ||
| return count > 0 | ||
| } | ||
|
|
||
| func TestLease(t *testing.T) { | ||
| ctx := log.ContextWithNewDefaultLogger(context.Background()) | ||
| conn := sqltest.OpenForTesting(ctx, t) | ||
| dal, err := New(ctx, conn) | ||
| assert.NoError(t, err) | ||
|
|
||
| _, err = dal.AcquireLease(ctx, leases.SystemKey("test"), time.Second*1) | ||
| assert.Error(t, err) | ||
|
|
||
| leasei, err := dal.AcquireLease(ctx, leases.SystemKey("test"), time.Second*5) | ||
| assert.NoError(t, err) | ||
|
|
||
| lease := leasei.(*Lease) //nolint:forcetypeassert | ||
|
|
||
| // Try to acquire the same lease again, which should fail. | ||
| _, err = dal.AcquireLease(ctx, leases.SystemKey("test"), time.Second*5) | ||
| assert.IsError(t, err, ErrConflict) | ||
|
|
||
| time.Sleep(time.Second * 6) | ||
|
|
||
| assert.True(t, leaseExists(t, conn, lease.idempotencyKey, lease.key)) | ||
|
|
||
| err = lease.Release() | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.False(t, leaseExists(t, conn, lease.idempotencyKey, lease.key)) | ||
| } | ||
|
|
||
| func TestExpireLeases(t *testing.T) { | ||
| ctx := log.ContextWithNewDefaultLogger(context.Background()) | ||
| conn := sqltest.OpenForTesting(ctx, t) | ||
| dal, err := New(ctx, conn) | ||
| assert.NoError(t, err) | ||
|
|
||
| leasei, err := dal.AcquireLease(ctx, leases.SystemKey("test"), time.Second*5) | ||
| assert.NoError(t, err) | ||
|
|
||
| lease := leasei.(*Lease) //nolint:forcetypeassert | ||
|
|
||
| err = dal.ExpireLeases(ctx) | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.True(t, leaseExists(t, conn, lease.idempotencyKey, lease.key)) | ||
|
|
||
| // Pretend that the lease expired. | ||
| lease.leak = true | ||
| err = lease.Release() | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.True(t, leaseExists(t, conn, lease.idempotencyKey, lease.key)) | ||
|
|
||
| time.Sleep(time.Second * 6) | ||
|
|
||
| err = dal.ExpireLeases(ctx) | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.False(t, leaseExists(t, conn, lease.idempotencyKey, lease.key)) | ||
|
|
||
| leasei, err = dal.AcquireLease(ctx, leases.SystemKey("test"), time.Second*5) | ||
| assert.NoError(t, err) | ||
|
|
||
| err = leasei.Release() | ||
| assert.NoError(t, err) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| package leases | ||
|
|
||
| import ( | ||
| "context" | ||
| "time" | ||
|
|
||
| "github.com/puzpuzpuz/xsync/v3" | ||
| ) | ||
|
|
||
| func NewFakeLeaser() *FakeLeaser { | ||
| return &FakeLeaser{ | ||
| leases: xsync.NewMapOf[string, struct{}](), | ||
| } | ||
| } | ||
|
|
||
| var _ Leaser = (*FakeLeaser)(nil) | ||
|
|
||
| // FakeLeaser is a fake implementation of the [Leaser] interface. | ||
| type FakeLeaser struct { | ||
| leases *xsync.MapOf[string, struct{}] | ||
| } | ||
|
|
||
| func (f *FakeLeaser) AcquireLease(ctx context.Context, key Key, ttl time.Duration) (Lease, error) { | ||
| if _, loaded := f.leases.LoadOrStore(key.String(), struct{}{}); loaded { | ||
| return nil, ErrConflict | ||
| } | ||
| return &FakeLease{leaser: f, key: key}, nil | ||
| } | ||
|
|
||
| type FakeLease struct { | ||
| leaser *FakeLeaser | ||
| key Key | ||
| } | ||
|
|
||
| func (f *FakeLease) Release() error { | ||
| f.leaser.leases.Delete(f.key.String()) | ||
| return nil | ||
| } | ||
|
|
||
| func (f *FakeLease) String() string { return f.key.String() } |
Oops, something went wrong.