优化结构

internal/integration/cache/cache.go

@@ -89,15 +89,23 @@ func (s *Store) Delete(ctx context.Context, key string) error {
 
 func (s *Store) Increment(ctx context.Context, key string, expiration time.Duration) (int64, error) {
 	if client := s.client(); client != nil {
+		if expiration > 0 {
+			expirationMS := expiration.Milliseconds()
+			if expirationMS < 1 {
+				expirationMS = 1
+			}
+			const incrementWithTTL = `
+local value = redis.call('INCR', KEYS[1])
+if redis.call('PTTL', KEYS[1]) < 0 then
+  redis.call('PEXPIRE', KEYS[1], ARGV[1])
+end
+return value`
+			return client.Eval(ctx, incrementWithTTL, []string{key}, expirationMS).Int64()
+		}
 		value, err := client.Incr(ctx, key).Result()
 		if err != nil {
 			return 0, err
 		}
-		if value == 1 && expiration > 0 {
-			if err := client.Expire(ctx, key, expiration).Err(); err != nil {
-				return 0, err
-			}
-		}
 		return value, nil
 	}
 	s.mu.Lock()
@@ -120,6 +128,8 @@ func (s *Store) Increment(ctx context.Context, key string, expiration time.Durat
 		} else {
 			entry.expiresAt = time.Time{}
 		}
+	} else if entry.expiresAt.IsZero() && expiration > 0 {
+		entry.expiresAt = time.Now().Add(expiration)
 	}
 	s.memory[key] = entry
 	return value, nil

internal/integration/cache/cache_test.go

@@ -40,3 +40,18 @@ func TestStoreIncrementRecreatesExpiredKeyWithoutStaleExpiry(t *testing.T) {
 		t.Fatalf("Get() after recreation = %q, %v, %v", value, ok, err)
 	}
 }
+
+func TestStoreIncrementAddsTTLToExistingPersistentCounter(t *testing.T) {
+	store := New(nil)
+	ctx := context.Background()
+	if err := store.Set(ctx, "counter", "5", 0); err != nil {
+		t.Fatal(err)
+	}
+	if value, err := store.Increment(ctx, "counter", time.Millisecond); err != nil || value != 6 {
+		t.Fatalf("Increment() = %d, %v", value, err)
+	}
+	time.Sleep(5 * time.Millisecond)
+	if _, ok, err := store.Get(ctx, "counter"); err != nil || ok {
+		t.Fatalf("persistent counter did not expire: ok=%v err=%v", ok, err)
+	}
+}