WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/common/redislock/token_lock.go

@@ -0,0 +1,324 @@
+// Package redislock provides token-owned Redis locks with atomic renewal and
+// release. Callers must treat loss of ownership as loss of exclusive access.
+package redislock
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const releaseTimeout = 5 * time.Second
+
+// ErrLockOwnershipLost means exclusive ownership can no longer be guaranteed.
+var ErrLockOwnershipLost = errors.New("redis lock ownership lost")
+
+type ownershipContextKey struct{}
+type ownershipCheckContextKey struct{}
+
+// EnsureOwnership is a final lease check before a durable mutation. A canceled
+// context or lost lease must never permit publishing a stale generated result.
+// Callers outside a renewable lock still honor normal context cancellation.
+func EnsureOwnership(ctx context.Context) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if check, ok := ctx.Value(ownershipCheckContextKey{}).(func() error); ok {
+		return check()
+	}
+	return nil
+}
+
+var (
+	releaseScript = redis.NewScript(`
+if redis.call('GET', KEYS[1]) == ARGV[1] then
+	return redis.call('DEL', KEYS[1])
+end
+return 0
+`)
+	renewScript = redis.NewScript(`
+if redis.call('GET', KEYS[1]) == ARGV[1] then
+	return redis.call('PEXPIRE', KEYS[1], ARGV[2])
+end
+return 0
+`)
+)
+
+// OwnershipContext returns a context that is canceled only when lock
+// ownership is lost, not when the request context is canceled. Outside a
+// renewable-lock callback it falls back to a cancellation-detached context.
+func OwnershipContext(ctx context.Context) context.Context {
+	if ctx == nil {
+		return context.Background()
+	}
+	if ownershipCtx, ok := ctx.Value(ownershipContextKey{}).(context.Context); ok {
+		return ownershipCtx
+	}
+	return context.WithoutCancel(ctx)
+}
+
+// NewToken returns a random owner token suitable for compare-and-* scripts.
+func NewToken() (string, error) {
+	var token [16]byte
+	if _, err := rand.Read(token[:]); err != nil {
+		return "", fmt.Errorf("generate redis lock token: %w", err)
+	}
+	return hex.EncodeToString(token[:]), nil
+}
+
+// TryAcquire attempts one SET NX acquisition.
+func TryAcquire(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	token string,
+	lease time.Duration,
+) (bool, error) {
+	if err := validateOperation(client, key, token, lease); err != nil {
+		return false, err
+	}
+	acquired, err := client.SetNX(ctx, key, token, lease).Result()
+	if err != nil {
+		return false, fmt.Errorf("acquire redis lock %q: %w", key, err)
+	}
+	return acquired, nil
+}
+
+// Acquire waits until the lock is acquired or ctx is cancelled.
+func Acquire(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	token string,
+	lease time.Duration,
+) error {
+	for {
+		acquired, err := TryAcquire(ctx, client, key, token, lease)
+		if err != nil {
+			return err
+		}
+		if acquired {
+			return nil
+		}
+
+		timer := time.NewTimer(retryDelay())
+		select {
+		case <-ctx.Done():
+			if !timer.Stop() {
+				<-timer.C
+			}
+			return ctx.Err()
+		case <-timer.C:
+		}
+	}
+}
+
+// Release deletes the lock only if token still owns it.
+func Release(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	token string,
+) (bool, error) {
+	if client == nil {
+		return false, errors.New("redis client is required")
+	}
+	if key == "" || token == "" {
+		return false, errors.New("redis lock key and token are required")
+	}
+	released, err := releaseScript.Run(ctx, client, []string{key}, token).Int64()
+	if err != nil {
+		return false, fmt.Errorf("release redis lock %q: %w", key, err)
+	}
+	return released != 0, nil
+}
+
+// Renew extends the lease only if token still owns the lock.
+func Renew(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	token string,
+	lease time.Duration,
+) (bool, error) {
+	if err := validateOperation(client, key, token, lease); err != nil {
+		return false, err
+	}
+	renewed, err := renewScript.Run(
+		ctx,
+		client,
+		[]string{key},
+		token,
+		lease.Milliseconds(),
+	).Int64()
+	if err != nil {
+		return false, fmt.Errorf("renew redis lock %q: %w", key, err)
+	}
+	return renewed != 0, nil
+}
+
+// WithRenewableLock acquires key, renews it while fn runs, and atomically
+// releases it on exit. Renewal failure cancels the context passed to fn.
+func WithRenewableLock(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	lease time.Duration,
+	renewInterval time.Duration,
+	fn func(context.Context) error,
+) (resultErr error) {
+	return WithRenewableLockWait(ctx, client, key, lease, renewInterval, 0, fn)
+}
+
+// WithRenewableLockWait bounds only acquisition time; a long callback keeps
+// the parent context and a renewable lease. Zero wait means no extra deadline.
+func WithRenewableLockWait(
+	ctx context.Context,
+	client redis.UniversalClient,
+	key string,
+	lease time.Duration,
+	renewInterval time.Duration,
+	wait time.Duration,
+	fn func(context.Context) error,
+) (resultErr error) {
+	if fn == nil {
+		return errors.New("redis lock callback is required")
+	}
+	if renewInterval <= 0 || renewInterval >= lease {
+		return errors.New("redis lock renewal interval must be positive and shorter than the lease")
+	}
+
+	token, err := NewToken()
+	if err != nil {
+		return err
+	}
+	acquireCtx := ctx
+	if wait > 0 {
+		var cancel context.CancelFunc
+		acquireCtx, cancel = context.WithTimeout(ctx, wait)
+		defer cancel()
+	}
+	if err := Acquire(acquireCtx, client, key, token, lease); err != nil {
+		return err
+	}
+
+	lockCtx, cancelLock := context.WithCancelCause(ctx)
+	ownershipCtx, cancelOwnership := context.WithCancelCause(context.Background())
+	lockCtx = context.WithValue(lockCtx, ownershipContextKey{}, ownershipCtx)
+	lockCtx = context.WithValue(lockCtx, ownershipCheckContextKey{}, func() error {
+		renewed, err := Renew(lockCtx, client, key, token, lease)
+		if err != nil || !renewed {
+			lost := fmt.Errorf("%w for %q", ErrLockOwnershipLost, key)
+			cancelOwnership(lost)
+			cancelLock(lost)
+			return lost
+		}
+		return nil
+	})
+	renewCtx, stopRenewal := context.WithCancel(context.WithoutCancel(ctx))
+	renewResult := make(chan error, 1)
+	go renewLoop(
+		renewCtx,
+		cancelLock,
+		cancelOwnership,
+		client,
+		key,
+		token,
+		lease,
+		renewInterval,
+		renewResult,
+	)
+
+	defer func() {
+		stopRenewal()
+		renewErr := <-renewResult
+		cancelLock(nil)
+		cancelOwnership(nil)
+
+		releaseCtx, cancelRelease := context.WithTimeout(context.Background(), releaseTimeout)
+		released, releaseErr := Release(releaseCtx, client, key, token)
+		cancelRelease()
+		if releaseErr == nil && !released && renewErr == nil {
+			releaseErr = fmt.Errorf("%w for %q", ErrLockOwnershipLost, key)
+		}
+
+		resultErr = errors.Join(resultErr, renewErr, ctx.Err(), releaseErr)
+	}()
+
+	resultErr = fn(lockCtx)
+	return resultErr
+}
+
+func renewLoop(
+	ctx context.Context,
+	cancelLock context.CancelCauseFunc,
+	cancelOwnership context.CancelCauseFunc,
+	client redis.UniversalClient,
+	key string,
+	token string,
+	lease time.Duration,
+	renewInterval time.Duration,
+	result chan<- error,
+) {
+	ticker := time.NewTicker(renewInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			result <- nil
+			return
+		case <-ticker.C:
+			renewed, err := Renew(ctx, client, key, token, lease)
+			if err != nil {
+				if errors.Is(err, context.Canceled) && ctx.Err() != nil {
+					result <- nil
+					return
+				}
+				ownershipErr := fmt.Errorf("%w: %w", ErrLockOwnershipLost, err)
+				cancelOwnership(ownershipErr)
+				cancelLock(ownershipErr)
+				result <- ownershipErr
+				return
+			}
+			if !renewed {
+				ownershipErr := fmt.Errorf("%w for %q", ErrLockOwnershipLost, key)
+				cancelOwnership(ownershipErr)
+				cancelLock(ownershipErr)
+				result <- ownershipErr
+				return
+			}
+		}
+	}
+}
+
+func validateOperation(
+	client redis.UniversalClient,
+	key string,
+	token string,
+	lease time.Duration,
+) error {
+	if client == nil {
+		return errors.New("redis client is required")
+	}
+	if key == "" || token == "" {
+		return errors.New("redis lock key and token are required")
+	}
+	if lease <= 0 {
+		return errors.New("redis lock lease must be positive")
+	}
+	return nil
+}
+
+func retryDelay() time.Duration {
+	var random [1]byte
+	if _, err := rand.Read(random[:]); err != nil {
+		return 50 * time.Millisecond
+	}
+	return 25*time.Millisecond + time.Duration(random[0])%50*time.Millisecond
+}

internal/common/redislock/token_lock_test.go

@@ -0,0 +1,269 @@
+package redislock_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+
+	"github.com/Tencent/WeKnora/internal/common/redislock"
+)
+
+func newTokenLockTestClient(t *testing.T) (*redis.Client, *miniredis.Miniredis) {
+	t.Helper()
+
+	mini := miniredis.RunT(t)
+	client := redis.NewClient(&redis.Options{Addr: mini.Addr()})
+	t.Cleanup(func() { require.NoError(t, client.Close()) })
+	return client, mini
+}
+
+func TestRenewableLockWaitOnlyLimitsAcquisition(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	err := redislock.WithRenewableLockWait(context.Background(), client, "lock:long", time.Second, 10*time.Millisecond, 15*time.Millisecond, func(ctx context.Context) error {
+		time.Sleep(50 * time.Millisecond)
+		require.NoError(t, ctx.Err(), "acquisition deadline must not cancel running synthesis")
+		return redislock.EnsureOwnership(ctx)
+	})
+	require.NoError(t, err)
+	require.NoError(t, client.Set(context.Background(), "lock:busy", "other-owner", time.Minute).Err())
+	called := false
+	err = redislock.WithRenewableLockWait(context.Background(), client, "lock:busy", time.Second, 10*time.Millisecond, 15*time.Millisecond, func(context.Context) error {
+		called = true
+		return nil
+	})
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.False(t, called)
+	require.Equal(t, "other-owner", client.Get(context.Background(), "lock:busy").Val())
+}
+
+func TestEnsureOwnershipRejectsStaleWriterBeforeNextRenewal(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	err := redislock.WithRenewableLockWait(context.Background(), client, "lock:stale", time.Minute, 20*time.Second, time.Second, func(ctx context.Context) error {
+		require.NoError(t, client.Set(context.Background(), "lock:stale", "successor", time.Minute).Err())
+		err := redislock.EnsureOwnership(ctx)
+		require.ErrorIs(t, err, redislock.ErrLockOwnershipLost)
+		require.Error(t, ctx.Err())
+		return err
+	})
+	require.ErrorIs(t, err, redislock.ErrLockOwnershipLost)
+	require.Equal(t, "successor", client.Get(context.Background(), "lock:stale").Val(), "stale release must not delete the successor")
+}
+
+func TestNewTokenReturnsUniqueOpaqueOwners(t *testing.T) {
+	t.Parallel()
+
+	first, err := redislock.NewToken()
+	require.NoError(t, err)
+	second, err := redislock.NewToken()
+	require.NoError(t, err)
+
+	require.Len(t, first, 32)
+	require.Len(t, second, 32)
+	require.NotEqual(t, first, second)
+}
+
+func TestTokenLockAcquireAndReleaseRequiresOwnership(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	ctx := context.Background()
+	key := "lock:ownership"
+
+	acquired, err := redislock.TryAcquire(ctx, client, key, "owner-a", time.Minute)
+	require.NoError(t, err)
+	require.True(t, acquired)
+
+	acquired, err = redislock.TryAcquire(ctx, client, key, "owner-b", time.Minute)
+	require.NoError(t, err)
+	require.False(t, acquired)
+
+	released, err := redislock.Release(ctx, client, key, "owner-b")
+	require.NoError(t, err)
+	require.False(t, released)
+	require.Equal(t, "owner-a", client.Get(ctx, key).Val())
+
+	released, err = redislock.Release(ctx, client, key, "owner-a")
+	require.NoError(t, err)
+	require.True(t, released)
+	require.ErrorIs(t, client.Get(ctx, key).Err(), redis.Nil)
+}
+
+func TestTokenLockRenewRequiresOwnership(t *testing.T) {
+	client, mini := newTokenLockTestClient(t)
+	ctx := context.Background()
+	key := "lock:renew"
+
+	require.NoError(t, client.Set(ctx, key, "owner-a", time.Second).Err())
+	renewed, err := redislock.Renew(ctx, client, key, "owner-b", time.Minute)
+	require.NoError(t, err)
+	require.False(t, renewed)
+
+	renewed, err = redislock.Renew(ctx, client, key, "owner-a", time.Minute)
+	require.NoError(t, err)
+	require.True(t, renewed)
+	mini.FastForward(2 * time.Second)
+	require.Equal(t, "owner-a", client.Get(ctx, key).Val())
+}
+
+func TestTokenLockRejectsUnsafeLeaseConfiguration(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	ctx := context.Background()
+
+	_, err := redislock.TryAcquire(ctx, client, "lock:key", "owner", 0)
+	require.Error(t, err)
+	_, err = redislock.Renew(ctx, client, "lock:key", "owner", -time.Second)
+	require.Error(t, err)
+	err = redislock.WithRenewableLock(
+		ctx,
+		client,
+		"lock:key",
+		time.Second,
+		time.Second,
+		func(context.Context) error { return nil },
+	)
+	require.Error(t, err)
+}
+
+func TestWithRenewableLockReturnsCallbackErrorAndReleases(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	ctx := context.Background()
+	key := "lock:callback"
+	want := errors.New("callback failed")
+
+	err := redislock.WithRenewableLock(
+		ctx,
+		client,
+		key,
+		time.Minute,
+		20*time.Second,
+		func(context.Context) error { return want },
+	)
+
+	require.ErrorIs(t, err, want)
+	require.ErrorIs(t, client.Get(ctx, key).Err(), redis.Nil)
+}
+
+func TestWithRenewableLockHonorsContextWhileWaiting(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	ctx := context.Background()
+	key := "lock:busy"
+	require.NoError(t, client.Set(ctx, key, "owner-a", time.Minute).Err())
+
+	waitCtx, cancel := context.WithTimeout(ctx, 25*time.Millisecond)
+	defer cancel()
+	called := false
+	err := redislock.WithRenewableLock(
+		waitCtx,
+		client,
+		key,
+		time.Minute,
+		20*time.Second,
+		func(context.Context) error {
+			called = true
+			return nil
+		},
+	)
+
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.False(t, called)
+}
+
+func TestWithRenewableLockCancelsCallbackAfterOwnershipLoss(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	key := "lock:lost"
+
+	err := redislock.WithRenewableLock(
+		context.Background(),
+		client,
+		key,
+		100*time.Millisecond,
+		10*time.Millisecond,
+		func(lockCtx context.Context) error {
+			require.NoError(t, client.Set(context.Background(), key, "replacement", time.Minute).Err())
+			select {
+			case <-lockCtx.Done():
+				return nil
+			case <-time.After(time.Second):
+				return errors.New("callback was not cancelled")
+			}
+		},
+	)
+
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "ownership lost")
+	require.Equal(t, "replacement", client.Get(context.Background(), key).Val())
+}
+
+func TestWithRenewableLockReportsOwnershipLossAtRelease(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	key := "lock:lost-before-renewal"
+
+	err := redislock.WithRenewableLock(
+		context.Background(),
+		client,
+		key,
+		time.Minute,
+		20*time.Second,
+		func(context.Context) error {
+			return client.Set(context.Background(), key, "replacement", time.Minute).Err()
+		},
+	)
+
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "ownership lost")
+	require.Equal(t, "replacement", client.Get(context.Background(), key).Val())
+}
+
+func TestOwnershipContextSurvivesCallerCancellation(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	err := redislock.WithRenewableLock(
+		ctx,
+		client,
+		"lock:caller-cancel",
+		time.Minute,
+		20*time.Second,
+		func(lockCtx context.Context) error {
+			ownershipCtx := redislock.OwnershipContext(lockCtx)
+			cancel()
+			require.ErrorIs(t, lockCtx.Err(), context.Canceled)
+			select {
+			case <-ownershipCtx.Done():
+				t.Fatal("caller cancellation must not imply lock ownership loss")
+			default:
+			}
+			return nil
+		},
+	)
+
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestOwnershipContextCancelsWhenRenewalLosesOwnership(t *testing.T) {
+	client, _ := newTokenLockTestClient(t)
+	key := "lock:ownership-context"
+
+	err := redislock.WithRenewableLock(
+		context.Background(),
+		client,
+		key,
+		100*time.Millisecond,
+		10*time.Millisecond,
+		func(lockCtx context.Context) error {
+			ownershipCtx := redislock.OwnershipContext(lockCtx)
+			require.NoError(t, client.Set(context.Background(), key, "replacement", time.Minute).Err())
+			select {
+			case <-ownershipCtx.Done():
+				return nil
+			case <-time.After(time.Second):
+				return errors.New("ownership context was not cancelled")
+			}
+		},
+	)
+
+	require.ErrorIs(t, err, redislock.ErrLockOwnershipLost)
+}