refactor: session locker

Author: chrisgacsal

pkg/framework/lockr/session.go

@@ -12,32 +12,63 @@ import (
 	"entgo.io/ent/dialect"
 	entsql "entgo.io/ent/dialect/sql"
 
+	"github.com/openmeterio/openmeter/pkg/clock"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 )
 
 var (
-	ErrNoLockAcquired         = errors.New("lock could not be acquired")
-	ErrNoLockReleased         = errors.New("lock could not be released")
-	ErrSessionLockerDone      = errors.New("session locker is already closed")
-	ErrDatabaseConnectionDown = errors.New("database connection is down")
+	ErrNoLockAcquired    = errors.New("lock could not be acquired")
+	ErrNoLockReleased    = errors.New("lock could not be released")
+	ErrSessionLockerDone = errors.New("session locker is already closed")
+	ErrSessionLockerBusy = errors.New("session locker is blocked by another lock request")
 )
 
 type Releaser func(context.Context) error
 
+type releaser struct {
+	done   atomic.Bool
+	locker *SessionLocker
+	key    Key
+}
+
+func (r *releaser) release(ctx context.Context) error {
+	if r.done.Load() {
+		return nil
+	}
+
+	rErr := r.locker.release(ctx, r.key)
+	if rErr != nil {
+		if !errors.Is(rErr, ErrNoLockReleased) && !errors.Is(rErr, ErrSessionLockerDone) {
+			return rErr
+		}
+	}
+
+	r.done.Store(true)
+
+	// Release references to locker and key so they can be GC'd
+	r.locker = nil
+	r.key = nil
+
+	return rErr
+}
+
 type SessionLockerConfig struct {
 	Logger         *slog.Logger
 	PostgresDriver *pgdriver.Driver
 }
 
+// SessionLocker is a locker that uses PostgreSQL advisory locks to acquire locks.
+// It requires a dedicated connection to acquire locks.
 type SessionLocker struct {
 	logger *slog.Logger
 	conn   *sql.Conn
 
 	closed atomic.Bool
 	closer func()
+	mu     sync.Mutex
 }
 
-func NewSessionLockr(config SessionLockerConfig) (*SessionLocker, error) {
+func NewSessionLockr(ctx context.Context, config SessionLockerConfig) (*SessionLocker, error) {
 	if config.Logger == nil {
 		return nil, errors.New("logger is required")
 	}
@@ -46,19 +77,23 @@ func NewSessionLockr(config SessionLockerConfig) (*SessionLocker, error) {
 		return nil, errors.New("postgres driver is required")
 	}
 
-	conn, err := config.PostgresDriver.DB().Conn(context.Background())
+	conn, err := config.PostgresDriver.DB().Conn(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get postgres connection: %w", err)
 	}
 
+	id := clock.Now().UTC().UnixNano()
+
+	logger := config.Logger.With("component", "session-lockr", "id", id)
+
 	closer := sync.OnceFunc(func() {
 		if err := conn.Close(); err != nil {
-			config.Logger.Error("failed to close postgres connection", "error", err)
+			logger.Error("failed to close postgres connection", "error", err)
 		}
 	})
 
 	return &SessionLocker{
-		logger: config.Logger,
+		logger: logger,
 		conn:   conn,
 		closer: closer,
 	}, nil
@@ -69,10 +104,6 @@ func (l *SessionLocker) lock(ctx context.Context, key Key, nonblocking bool) (Re
 		return nil, ErrSessionLockerDone
 	}
 
-	if err := l.conn.PingContext(ctx); err != nil {
-		return nil, ErrDatabaseConnectionDown
-	}
-
 	lockFunc := "pg_advisory_lock"
 
 	if nonblocking {
@@ -101,19 +132,17 @@ func (l *SessionLocker) lock(ctx context.Context, key Key, nonblocking bool) (Re
 		return nil, fmt.Errorf("failed to acquire session-level advisory lock: %w", checkForTimeout(err))
 	}
 
-	if nonblocking {
-		var locked bool
+	var lockAcquired bool
 
+	if nonblocking {
 		for rows.Next() {
-			if err := rows.Scan(&locked); err != nil {
+			if err := rows.Scan(&lockAcquired); err != nil {
 				return nil, fmt.Errorf("failed to scan session-level advisory lock result: %w", err)
 			}
 		}
-
-		if !locked {
-			return nil, ErrNoLockAcquired
-		}
 	} else {
+		lockAcquired = true
+
 		for rows.Next() {
 		}
 	}
@@ -122,30 +151,48 @@ func (l *SessionLocker) lock(ctx context.Context, key Key, nonblocking bool) (Re
 		return nil, checkForTimeout(err)
 	}
 
-	r := &struct {
-		once sync.Once
-	}{}
-
-	return func(rCtx context.Context) error {
-		var err error
+	if !lockAcquired {
+		return nil, ErrNoLockAcquired
+	}
 
-		r.once.Do(func() {
-			err = l.Release(rCtx, key)
-		})
+	r := &releaser{
+		locker: l,
+		key:    key,
+	}
 
-		return err
-	}, nil
+	return r.release, nil
 }
 
+// TryLock attempts to acquire a lock for the given key in a non-blocking way and returns a Releaser that can be used
+// to release the lock if it is successfully acquired. The ErrNoLockAcquired is acquiring the lock is denied by the database server.
+// It may return ErrSessionLockerBusy if the SessionLocker is blocked by another caller, indicating that the lock request may be retried.
+// The ErrSessionLockerDone is returned if SessionLocker is closed, meaning it cannot be used for acquiring locks.
 func (l *SessionLocker) TryLock(ctx context.Context, key Key) (Releaser, error) {
+	mutexLocked := l.mu.TryLock()
+	if !mutexLocked {
+		return nil, ErrSessionLockerBusy
+	}
+
+	defer l.mu.Unlock()
+
 	return l.lock(ctx, key, true)
 }
 
+// Lock blocks until a lock is acquired and returns a Releaser that can be used to release the lock if it is successfully acquired.
+// The ErrNoLockAcquired is acquiring the lock is denied by the database server.
+// The ErrSessionLockerDone is returned if SessionLocker is closed, meaning it cannot be used for acquiring locks.
 func (l *SessionLocker) Lock(ctx context.Context, key Key) (Releaser, error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
 	return l.lock(ctx, key, false)
 }
 
-func (l *SessionLocker) Release(ctx context.Context, key Key) error {
+func (l *SessionLocker) release(ctx context.Context, key Key) error {
+	if l.closed.Load() {
+		return ErrSessionLockerDone
+	}
+
 	q, args := entsql.Dialect(dialect.Postgres).
 		SelectExpr(entsql.ExprFunc(func(b *entsql.Builder) {
 			b.WriteString("pg_advisory_unlock")
@@ -168,10 +215,10 @@ func (l *SessionLocker) Release(ctx context.Context, key Key) error {
 		return fmt.Errorf("failed to release session-level advisory lock: %w", checkForTimeout(err))
 	}
 
-	var released bool
+	var lockReleased bool
 
 	for rows.Next() {
-		if err = rows.Scan(&released); err != nil {
+		if err = rows.Scan(&lockReleased); err != nil {
 			return fmt.Errorf("failed to scan session-level advisory lock release result: %w", err)
 		}
 	}
@@ -180,10 +227,25 @@ func (l *SessionLocker) Release(ctx context.Context, key Key) error {
 		return checkForTimeout(err)
 	}
 
+	if !lockReleased {
+		return ErrNoLockReleased
+	}
+
 	return nil
 }
 
+// Close releases all locks held by the SessionLocker and closes the underlying database connection.
 func (l *SessionLocker) Close() {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if l.closed.Load() {
+		return
+	}
+
 	l.closer()
 	l.closed.Store(true)
+
+	// Release references to conn so it can be GC'd
+	l.conn = nil
 }

pkg/framework/lockr/session_test.go

@@ -27,7 +27,7 @@ func newTestSessionLocker(t *testing.T, dbConn string, opts ...pgdriver.Option)
 		}
 	})
 
-	locker, err := NewSessionLockr(SessionLockerConfig{
+	locker, err := NewSessionLockr(t.Context(), SessionLockerConfig{
 		Logger:         testutils.NewLogger(t),
 		PostgresDriver: postgresDriver,
 	})
@@ -205,18 +205,6 @@ func Test_SessionLocker(t *testing.T) {
 		require.Equal(t, []string{"s2 waiting", "s1 releasing", "s2 acquired"}, results)
 	})
 
-	t.Run("Release without lock is a no-op", func(t *testing.T) {
-		locker := newTestSessionLocker(t, testDB.URL)
-		defer locker.Close()
-
-		key, err := NewKey("test", "release-noop")
-		require.NoError(t, err)
-
-		// Releasing a lock that was never acquired should not error
-		err = locker.Release(t.Context(), key)
-		require.NoError(t, err)
-	})
-
 	t.Run("Lock respects context cancellation", func(t *testing.T) {
 		locker1 := newTestSessionLocker(t, testDB.URL)
 		defer locker1.Close()