Khóa phân tán trong hệ thống phân tán: Redlock, PostgreSQL Advisory Locks & etcd Leases

// Simplified Redlock client (conceptual, not production-ready without fencing)
import { Redis } from 'ioredis';
import { v4 as uuidv4 } from 'uuid';

interface RedlockOptions {
  lockKey: string;
  resourceId: string;
  ttlMs: number; // Time-to-live for the lock in milliseconds
  retryDelayMs?: number;
  maxRetries?: number;
}

class RedlockClient {
  private redisClients: Redis[];
  private quorum: number;

  constructor(redisUrls: string[]) {
    this.redisClients = redisUrls.map(url => new Redis(url));
    this.quorum = Math.floor(redisUrls.length / 2) + 1;
    if (this.quorum === 0) {
      throw new Error("Redlock requires at least one Redis instance.");
    }
  }

  /**
   * Attempts to acquire a distributed lock using the Redlock algorithm.
   * @returns A unique token if the lock is acquired, otherwise null.
   */
  public async acquireLock(options: RedlockOptions): Promise<string | null> {
    const { lockKey, resourceId, ttlMs, retryDelayMs = 100, maxRetries = 5 } = options;
    const value = uuidv4(); // Unique token for this lock attempt
    let retries = 0;

    while (retries < maxRetries) {
      const startTime = Date.now();
      let acquiredCount = 0;

      const acquirePromises = this.redisClients.map(async (client) => {
        try {
          // SET key value NX PX ttlMs
          // NX: Only set if the key does not already exist.
          // PX: Set the specified expire time, in milliseconds.
          const result = await client.set(lockKey, value, 'NX', 'PX', ttlMs);
          return result === 'OK';
        } catch (error) {
          console.error(`Redis client error during acquire: ${error}`);
          return false;
        }
      });

      const results = await Promise.all(acquirePromises);
      acquiredCount = results.filter(Boolean).length;

      const elapsedTime = Date.now() - startTime;
      const isValid = elapsedTime < ttlMs;

      if (acquiredCount >= this.quorum && isValid) {
        console.log(`Lock '${lockKey}' acquired by '${value}' on ${acquiredCount} instances.`);
        return value; // Lock acquired successfully
      } else {
        // If not acquired, or validity time expired, release any locks we did get
        await this.releaseLock(lockKey, value);
        console.warn(`Failed to acquire lock '${lockKey}'. Acquired ${acquiredCount}/${this.quorum} instances. Retrying...`);
        retries++;
        await new Promise(resolve => setTimeout(resolve, retryDelayMs));
      }
    }

    console.error(`Failed to acquire lock '${lockKey}' after ${maxRetries} retries.`);
    return null;
  }

  /**
   * Releases a distributed lock.
   * It's crucial to only release locks that match the client's unique value.
   */
  public async releaseLock(lockKey: string, value: string): Promise<void> {
    // Lua script to ensure atomic check-and-delete
    const luaScript = `
      if redis.call("get", KEYS[1]) == ARGV[1] then
        return redis.call("del", KEYS[1])
      else
        return 0
      end
    `;

    const releasePromises = this.redisClients.map(async (client) => {
      try {
        await client.eval(luaScript, 1, lockKey, value);
      } catch (error) {
        console.error(`Redis client error during release: ${error}`);
      }
    });

    await Promise.all(releasePromises);
    console.log(`Lock '${lockKey}' released by '${value}'.`);
  }

  public async disconnect(): Promise<void> {
    await Promise.all(this.redisClients.map(client => client.quit()));
  }
}

// Example Usage (requires running Redis instances)
async function runRedlockExample() {
  const redisUrls = ['redis://localhost:6379', 'redis://localhost:6380', 'redis://localhost:6381'];
  const redlock = new RedlockClient(redisUrls);

  const lockKey = 'my_critical_resource';
  const resourceId = 'unique_resource_id_123';
  const ttlMs = 5000; // 5 seconds

  console.log('Attempting to acquire lock 1...');
  const token1 = await redlock.acquireLock({ lockKey, resourceId, ttlMs });

  if (token1) {
    console.log(`Client 1 acquired lock with token: ${token1}. Performing critical operation...`);
    // Simulate work
    await new Promise(resolve => setTimeout(resolve, 2000));
    console.log('Client 1 finished critical operation. Releasing lock...');
    await redlock.releaseLock(lockKey, token1);
  } else {
    console.log('Client 1 failed to acquire lock.');
  }

  console.log('\nAttempting to acquire lock 2 (should succeed after lock 1 is released)...');
  const token2 = await redlock.acquireLock({ lockKey, resourceId, ttlMs });

  if (token2) {
    console.log(`Client 2 acquired lock with token: ${token2}. Performing critical operation...`);
    await new Promise(resolve => setTimeout(resolve, 1000));
    await redlock.releaseLock(lockKey, token2);
  } else {
    console.log('Client 2 failed to acquire lock.');
  }

  await redlock.disconnect();
}

// To run this, you'd need 3 Redis instances running, e.g.:
// docker run -p 6379:6379 --name redis1 -d redis
// docker run -p 6380:6379 --name redis2 -d redis
// docker run -p 6381:6379 --name redis3 -d redis
// Then: node your_script.js
// runRedlockExample();

import { Pool, PoolClient } from 'pg';

interface AdvisoryLockOptions {
  lockId: number; // A unique BIGINT identifier for the lock
  timeoutMs?: number; // How long to wait for the lock (0 for non-blocking)
}

class PgAdvisoryLock {
  private pool: Pool;

  constructor(connectionString: string) {
    this.pool = new Pool({ connectionString });
  }

  /**
   * Acquires a transaction-level advisory lock.
   * The lock is automatically released when the transaction commits or rolls back.
   * @returns The PoolClient if the lock is acquired, otherwise null.
   */
  public async acquireTransactionLock(options: AdvisoryLockOptions): Promise<PoolClient | null> {
    const { lockId, timeoutMs = 0 } = options;
    const client = await this.pool.connect();

    try {
      await client.query('BEGIN'); // Start a transaction

      let lockAcquired = false;
      if (timeoutMs === 0) {
        // Non-blocking attempt
        const res = await client.query('SELECT pg_try_advisory_xact_lock($1)', [lockId]);
        lockAcquired = res.rows[0].pg_try_advisory_xact_lock;
      } else {
        // Blocking attempt with timeout
        // pg_advisory_xact_lock will block until acquired or connection reset.
        // We simulate a timeout by running it in a separate promise and racing.
        const acquirePromise = client.query('SELECT pg_advisory_xact_lock($1)', [lockId]);
        const timeoutPromise = new Promise<void>(resolve => setTimeout(() => resolve(), timeoutMs));

        await Promise.race([acquirePromise, timeoutPromise]);
        // If acquirePromise resolved, lockAcquired is true. If timeoutPromise resolved first, it's false.
        // This is a simplification; a more robust solution might involve a separate connection for the timeout check.
        // For pg_advisory_xact_lock, it either succeeds or blocks indefinitely.
        // pg_try_advisory_xact_lock is generally preferred for non-blocking or explicit retry logic.
        // For a true blocking with timeout, one might need a loop with pg_try_advisory_xact_lock.
        const res = await client.query('SELECT pg_try_advisory_xact_lock($1)', [lockId]); // Re-check after potential block
        lockAcquired = res.rows[0].pg_try_advisory_xact_lock;
      }

      if (lockAcquired) {
        console.log(`Transaction lock ${lockId} acquired.`);
        return client; // Return the client to the caller for transaction operations
      } else {
        await client.query('ROLLBACK'); // Rollback if lock not acquired
        client.release();
        console.warn(`Failed to acquire transaction lock ${lockId}.`);
        return null;
      }
    } catch (error) {
      await client.query('ROLLBACK');
      client.release();
      console.error(`Error acquiring transaction lock ${lockId}: ${error}`);
      throw error;
    }
  }

  /**
   * Releases a transaction-level advisory lock by committing or rolling back the transaction.
   * This method should be called by the client that acquired the lock.
   */
  public async releaseTransactionLock(client: PoolClient, success: boolean): Promise<void> {
    try {
      if (success) {
        await client.query('COMMIT');
        console.log('Transaction committed, lock released.');
      } else {
        await client.query('ROLLBACK');
        console.log('Transaction rolled back, lock released.');
      }
    } catch (error) {
      console.error(`Error releasing transaction lock: ${error}`);
      await client.query('ROLLBACK'); // Ensure rollback on error
    } finally {
      client.release();
    }
  }

  /**
   * Acquires a session-level advisory lock.
   * The lock persists until explicitly released or the session ends.
   * @returns true if lock is acquired, false otherwise.
   */
  public async acquireSessionLock(options: AdvisoryLockOptions): Promise<boolean> {
    const { lockId, timeoutMs = 0 } = options;
    const client = await this.pool.connect(); // New client for session lock

    try {
      let lockAcquired = false;
      if (timeoutMs === 0) {
        const res = await client.query('SELECT pg_try_advisory_lock($1)', [lockId]);
        lockAcquired = res.rows[0].pg_try_advisory_lock;
      } else {
        // For session locks, pg_advisory_lock blocks.
        // A timeout would require a separate mechanism or loop with pg_try_advisory_lock.
        // For simplicity, we'll use pg_try_advisory_lock with retries for timeout simulation.
        const startTime = Date.now();
        while (Date.now() - startTime < timeoutMs) {
          const res = await client.query('SELECT pg_try_advisory_lock($1)', [lockId]);
          if (res.rows[0].pg_try_advisory_lock) {
            lockAcquired = true;
            break;
          }
          await new Promise(resolve => setTimeout(resolve, 50)); // Small delay before retry
        }
      }

      if (lockAcquired) {
        console.log(`Session lock ${lockId} acquired.`);
        // Store the client if you need to explicitly release it later
        // For this example, we'll just return true and assume the caller manages the client.
        // In a real app, you'd likely keep a map of lockId -> client.
        return true;
      } else {
        client.release(); // Release client if lock not acquired
        console.warn(`Failed to acquire session lock ${lockId}.`);
        return false;
      }
    } catch (error) {
      client.release();
      console.error(`Error acquiring session lock ${lockId}: ${error}`);
      throw error;
    }
  }

  /**
   * Releases a session-level advisory lock.
   * Requires the same client that acquired the lock.
   */
  public async releaseSessionLock(lockId: number, client: PoolClient): Promise<void> {
    try {
      await client.query('SELECT pg_advisory_unlock($1)', [lockId]);
      console.log(`Session lock ${lockId} released.`);
    } catch (error) {
      console.error(`Error releasing session lock ${lockId}: ${error}`);
    } finally {
      client.release();
    }
  }

  public async disconnect(): Promise<void> {
    await this.pool.end();
  }
}

// Example Usage (requires running PostgreSQL)
async function runPgAdvisoryLockExample() {
  const connectionString = 'postgresql://user:password@localhost:5432/mydatabase';
  const pgLock = new PgAdvisoryLock(connectionString);
  const lockId = 12345;

  // --- Transaction-level lock example ---
  console.log('\n--- Transaction-level lock ---');
  let client1: PoolClient | null = null;
  try {
    client1 = await pgLock.acquireTransactionLock({ lockId, timeoutMs: 100 });
    if (client1) {
      console.log('Client 1 acquired transaction lock. Performing database operations...');
      // Simulate database operation within the transaction
      await client1.query('INSERT INTO some_table (data) VALUES ($1)', ['transactional_data_1']);
      await new Promise(resolve => setTimeout(resolve, 1000)); // Simulate work
      await pgLock.releaseTransactionLock(client1, true); // Commit and release
    } else {
      console.log('Client 1 failed to acquire transaction lock.');
    }
  } catch (error) {
    console.error('Transaction example error:', error);
    if (client1) await pgLock.releaseTransactionLock(client1, false); // Rollback
  }

  // --- Session-level lock example ---
  console.log('\n--- Session-level lock ---');
  let sessionClient: PoolClient | null = null;
  try {
    sessionClient = await pgLock.pool.connect(); // Get a client for the session lock
    const acquired = await pgLock.acquireSessionLock({ lockId: lockId + 1, timeoutMs: 100 }); // Use a different lock ID

    if (acquired) {
      console.log('Client 1 acquired session lock. Performing operations...');
      // Simulate work
      await new Promise(resolve => setTimeout(resolve, 1500));

      // Attempt to acquire the same session lock from another client (should fail)
      console.log('Client 2 attempting to acquire the same session lock...');
      const client2 = await pgLock.pool.connect();
      const acquired2 = await pgLock.acquireSessionLock({ lockId: lockId + 1, timeoutMs: 100 });
      if (!acquired2) {
        console.log('Client 2 correctly failed to acquire the session lock.');
      }
      client2.release(); // Release client2 connection

      await pgLock.releaseSessionLock(lockId + 1, sessionClient);
    } else {
      console.log('Client 1 failed to acquire session lock.');
      if (sessionClient) sessionClient.release();
    }
  } catch (error) {
    console.error('Session example error:', error);
    if (sessionClient) sessionClient.release();
  }

  await pgLock.disconnect();
}

// To run this, you'd need a PostgreSQL instance running, e.g.:
// docker run -p 5432:5432 --name pg-lock -e POSTGRES_USER=user -e POSTGRES_PASSWORD=password -e POSTGRES_DB=mydatabase -d postgres
// Then: node your_script.js
// runPgAdvisoryLockExample();

import { Etcd3, IOptions } from 'etcd3';

interface EtcdLockOptions {
  lockKey: string;
  ttlSeconds: number; // Lease time-to-live in seconds
  retryDelayMs?: number;
  maxRetries?: number;
}

class EtcdDistributedLock {
  private etcd: Etcd3;

  constructor(etcdHosts: string | string[], options?: IOptions) {
    this.etcd = new Etcd3({ hosts: etcdHosts, ...options });
  }

  /**
   * Acquires a distributed lock using etcd leases and compare-and-swap.
   * Returns the lease ID and the revision number (fencing token) if successful.
   */
  public async acquireLock(options: EtcdLockOptions): Promise<{ leaseId: string; revision: number } | null> {
    const { lockKey, ttlSeconds, retryDelayMs = 100, maxRetries = 10 } = options;
    const clientValue = Math.random().toString(36).substring(2, 15); // Unique client identifier

    let retries = 0;
    while (retries < maxRetries) {
      try {
        // 1. Create a lease
        const lease = this.etcd.lease(ttlSeconds);
        const leaseId = lease.id;

        // 2. Attempt to acquire the lock using a transaction (compare-and-swap)
        // Ensure the key does not exist, then create it with our value and lease.
        const transactionResult = await this.etcd.txn()
          .if(this.etcd.op(lockKey).not.exists())
          .then(this.etcd.op(lockKey).put(clientValue).withLease(leaseId))
          .commit();

        if (transactionResult.succeeded) {
          // Lock acquired. Start keep-alive.
          lease.on('lost', () => {
            console.error(`Lease ${leaseId} for lock ${lockKey} lost!`);
            // In a real application, this would trigger recovery or error handling.
          });
          lease.on('expired', () => {
            console.warn(`Lease ${leaseId} for lock ${lockKey} expired!`);
          });
          lease.on('keepalive established', () => {
            // console.log(`Keep-alive established for lease ${leaseId}`);
          });
          lease.on('keepalive error', (err) => {
            console.error(`Keep-alive error for lease ${leaseId}: ${err}`);
          });

          // Get the revision number for fencing
          const getResult = await this.etcd.get(lockKey).string();
          if (getResult) {
            const revision = getResult.modRevision;
            console.log(`Lock '${lockKey}' acquired by client '${clientValue}' with lease ${leaseId}, revision ${revision}`);
            return { leaseId, revision };
          } else {
            // This should ideally not happen if transaction succeeded, but defensive check.
            console.error(`Failed to get revision for lock ${lockKey} after acquiring.`);
            await this.releaseLock(lockKey, leaseId); // Clean up
            return null;
          }
        } else {
          // Lock already held by another client
          console.warn(`Lock '${lockKey}' already held. Retrying...`);
          await lease.revoke(); // Revoke our unused lease
          retries++;
          await new Promise(resolve => setTimeout(resolve, retryDelayMs));
        }
      } catch (error) {
        console.error(`Error acquiring lock '${lockKey}': ${error}`);
        retries++;
        await new Promise(resolve => setTimeout(resolve, retryDelayMs));
      }
    }

    console.error(`Failed to acquire lock '${lockKey}' after ${maxRetries} retries.`);
    return null;
  }

  /**
   * Releases a distributed lock by revoking its lease.
   * This will automatically delete the associated key.
   */
  public async releaseLock(lockKey: string, leaseId: string): Promise<void> {
    try {
      // Revoking the lease automatically deletes all keys associated with it.
      await this.etcd.lease.revoke(leaseId);
      console.log(`Lock '${lockKey}' with lease ${leaseId} released.`);
    } catch (error) {
      console.error(`Error releasing lock '${lockKey}' with lease ${leaseId}: ${error}`);
    }
  }

  /**
   * Verifies the fencing token (revision) for an operation.
   * The resource being protected should call this before performing an operation.
   */
  public async verifyFencingToken(lockKey: string, expectedRevision: number): Promise<boolean> {
    try {
      const getResult = await this.etcd.get(lockKey).string();
      if (!getResult) {
        console.warn(`Fencing check failed: Lock key '${lockKey}' does not exist.`);
        return false;
      }
      if (getResult.modRevision !== expectedRevision) {
        console.warn(`Fencing check failed: Expected revision ${expectedRevision}, got ${getResult.modRevision}.`);
        return false;
      }
      return true;
    } catch (error) {
      console.error(`Error during fencing token verification for '${lockKey}': ${error}`);
      return false;
    }
  }

  public async disconnect(): Promise<void> {
    await this.etcd.close();
  }
}

// Example Usage (requires running etcd)
async function runEtcdLockExample() {
  const etcdHosts = ['localhost:2379']; // Or an array of hosts for a cluster
  const etcdLock = new EtcdDistributedLock(etcdHosts);

  const lockKey = '/app/resource/processor_lock';
  const ttlSeconds = 5; // Lease expires in 5 seconds if not kept alive

  console.log('Attempting to acquire lock 1...');
  const lockInfo1 = await etcdLock.acquireLock({ lockKey, ttlSeconds });

  if (lockInfo1) {
    console.log(`Client 1 acquired lock with lease ${lockInfo1.leaseId}, revision ${lockInfo1.revision}.`);
    // Simulate critical operation with fencing check
    const isFenced = await etcdLock.verifyFencingToken(lockKey, lockInfo1.revision);
    if (isFenced) {
      console.log('Fencing token verified. Performing critical operation...');
      await new Promise(resolve => setTimeout(resolve, 3000)); // Simulate work
      console.log('Client 1 finished critical operation. Releasing lock...');
      await etcdLock.releaseLock(lockKey, lockInfo1.leaseId);
    } else {
      console.error('Fencing check failed for client 1. Aborting operation.');
      await etcdLock.releaseLock(lockKey, lockInfo1.leaseId); // Ensure release
    }
  } else {
    console.log('Client 1 failed to acquire lock.');
  }

  console.log('\nAttempting to acquire lock 2 (should succeed after lock 1 is released)...');
  const lockInfo2 = await etcdLock.acquireLock({ lockKey, ttlSeconds });

  if (lockInfo2) {
    console.log(`Client 2 acquired lock with lease ${lockInfo2.leaseId}, revision ${lockInfo2.revision}.`);
    await new Promise(resolve => setTimeout(resolve, 1000));
    await etcdLock.releaseLock(lockKey, lockInfo2.leaseId);
  } else {
    console.log('Client 2 failed to acquire lock.');
  }

  await etcdLock.disconnect();
}

// To run this, you'd need an etcd instance running, e.g.:
// docker run -p 2379:2379 -p 2380:2380 --name etcd-single -d quay.io/coreos/etcd:latest etcd -advertise-client-urls http://0.0.0.0:2379 -listen-client-urls http://0.0.0.0:2379
// Then: node your_script.js
// runEtcdLockExample();