JAVA之AQS

前言

由于去年12月底刚换工作，再加上今年3月份上海开始封控；一直处于高度紧张的状态下，没法静下心来好好geek。

AQS这个主题也是老生常谈了，面试中也经常遇到；但大多资料都是理论上的，各个属性是什么意思以及源码分析，工作中也是使用AbstractQueuedSynchronizer的实现类，所以想通过自己实现一个AQS来加深对其的把握。

这里想实现一个多key的ReentrantLock，支持以key进行加/释放锁。效果可以参考redis的setnx, 区别是可以自动唤醒，setnx可能还需要配合自旋。

PS: 下面实现的只是一个demo，只是为了加深对AQS的理解。

实现

STEP1

要制作自己的AQS首先要继承AbstractQueuedSynchronizer并且实现以下几个方法：

public class MultiReentrantLock extends AbstractQueuedSynchronizer {
    ...
    // 是否可以拿到锁，false的话会进入队列，当前Thread也会挂起
    protected boolean tryAcquire(int arg) {
        return super.tryAcquire(arg);
    }

    // 是否可以释放，true的话激活第一个等待的节点
    protected boolean tryRelease(int arg) {
        return super.tryRelease(arg);
    }

    // 共享锁
    protected int tryAcquireShared(int arg) {
        return super.tryAcquireShared(arg);
    }

    // 释放共享锁
    protected boolean tryReleaseShared(int arg) {
        return super.tryReleaseShared(arg);
    }
}

STEP2

注意到tryAcquire与tryRelease只有一个int类型的入参，并不支持我们传递一个key进来进行判断。ThreadLocal是一种实现方式，在tryAcquire之前设置，在tryRelease中获取，并进行一些逻辑判断，最后在release之后删除；但是在编写中发现，ThreadLocal不能满足我的需求(与STEP3提到的问题有关)，因为在tryRelease阶段，需要在A线程的上下文去获取跟B线程绑定的key。所以在实际的demo里，是将key设置到thread name中的。

public class MultiReentrantLock extends AbstractQueuedSynchronizer {
    ...
    public void lock(String lockKey) {
        Thread.currentThread().setName(lockKey);
        acquire(1);
    }
    ...
}

STEP3

这里会有一个问题，我们重温下acquire与release的逻辑：

acquire

    public final void acquire(int arg) {
        // 1. 判断是否可以拿到锁
        // 2. 如果不可以的话，将当前节点加入队列进行等待
        if (!tryAcquire(arg) &&
            acquireQueued(addWaiter(Node.EXCLUSIVE), arg))
            selfInterrupt();
    }
    
    // 创建节点并将其添加到队列尾部
    private Node addWaiter(Node mode) {
        Node node = new Node(Thread.currentThread(), mode);
        Node pred = tail;
        if (pred != null) {
            node.prev = pred;
            if (compareAndSetTail(pred, node)) {
                pred.next = node;
                return node;
            }
        }
        enq(node);
        return node;
    }
    
    final boolean acquireQueued(final Node node, int arg) {
        boolean failed = true;
        try {
            boolean interrupted = false;
            for (;;) {
                final Node p = node.predecessor();
                // 如果前节点是head节点，并且当前可以拿到锁就将自己设置为头节点
                if (p == head && tryAcquire(arg)) {
                    setHead(node);
                    p.next = null; // help GC
                    failed = false;
                    return interrupted;
                }
                // 如果还需要等待，先将节点状态设置为`SIGNAL`(等待信息)
                // 然后再将当前现场挂起
                // 当线程重新可运行时，判断其状态是否为interrupted，是的话会执行`selfInterrupt`
                if (shouldParkAfterFailedAcquire(p, node) &&
                    parkAndCheckInterrupt())
                    interrupted = true;
            }
        } finally {
            // 如果有异常失败了，则将该节点从队列中移除
            if (failed)
                cancelAcquire(node);
        }
    }

release

    public final boolean release(int arg) {
        // 如果可以释放，则激活下一个节点
        if (tryRelease(arg)) {
            Node h = head;
            if (h != null && h.waitStatus != 0)
                unparkSuccessor(h);
            return true;
        }
        return false;
    }
    
    private void unparkSuccessor(Node node) {
        int ws = node.waitStatus;
        if (ws < 0)
            compareAndSetWaitStatus(node, ws, 0);

        Node s = node.next;
        if (s == null || s.waitStatus > 0) {
            s = null;
            for (Node t = tail; t != null && t != node; t = t.prev)
                if (t.waitStatus <= 0)
                    s = t;
        }
        if (s != null)
            LockSupport.unpark(s.thread);
    }
    
    // 这里LockSupport.unpark(s.thread)后就会转到acquireQueued并执行以下逻辑
    final boolean acquireQueued(final Node node, int arg) {
        boolean failed = true;
        try {
            boolean interrupted = false;
            for (;;) {
                final Node p = node.predecessor();
                if (p == head && tryAcquire(arg)) {
                    setHead(node);
                    p.next = null; // help GC
                    failed = false;
                    return interrupted;
                }
                ...
            }
        } finally {
            ...
        }
    }

问题

STEP3开头提到有一个问题，在这里说明下：

在acquire需要加入队列时，添加到队列的节点是无key信息的；release也一样，只是激活下一个节点的线程，并不是激活与当前节点有相同key的节点。所以需要在tryAcquire与tryRelease中有一些逻辑判断：

1. 在tryAcquire中，我们需要判断是否有相同的key的thread在执行中或者在等待中？如果没有，就直接执行；否则需要加入等待队列；
2. 在tryRelease中，我们需要判断下一个待激活的节点绑定的key释放跟当前要释放的一样？如果一样，就返回true激活下一个节点；否则，就需要再判断下一个待激活的节点所关联的key是否有其他线程正在执行，如果没有就激活；否则不激活。

所以在这里，我们再将key所关联的所有Thread保存在一个数组里。

public class MultiReentrantLock extends AbstractQueuedSynchronizer {

    ConcurrentHashMap<String, LinkedList<Thread>> lockQueueHolders = new ConcurrentHashMap<>();
}

完整代码展示

public class MultiReentrantLock extends AbstractQueuedSynchronizer {

    ConcurrentHashMap<String, LinkedList<Thread>> lockQueueHolders = new ConcurrentHashMap<>();

    public MultiReentrantLock() {
    }

    public void lock(String lockKey) {
        Thread.currentThread().setName(lockKey);
        acquire(1);
    }

    public void unlock(String lockKey) {
        if (null == lockKey) {
            return;
        }
        // 要释放的key与自身关联的不一致，抛异常
        if (!lockKey.equals(Thread.currentThread().getName())) {
            throw new UnsupportedOperationException();
        }
        release(1);
    }

    @Override
    protected boolean tryAcquire(int arg) {
        LinkedList<Thread> threadQueue = lockQueueHolders.getOrDefault(Thread.currentThread().getName(), new LinkedList<>());
        boolean ret = false;
        // 如果该key还未设置threadQueue或者第一个thread与当前thread相同，暂时判定结果为true
        if (threadQueue.isEmpty()) {
            ret = true;
            addToLockQueue(threadQueue);
        } else if (threadQueue.getFirst() == Thread.currentThread()) {
            ret = true;
        } else {
            addToLockQueue(threadQueue);
        }
        LinkedList<Thread> curThreadQueue = lockQueueHolders.putIfAbsent(Thread.currentThread().getName(), threadQueue);
        // 如果有别的线程提早将相同key对应的threadQueue设置到Map中，则判定为失败
        // 同时将其加入到已有的threadQueue中
        if (null != curThreadQueue && curThreadQueue != threadQueue) {
            addToLockQueue(curThreadQueue);
            ret = false;
        }
        return ret;
    }

    private void addToLockQueue(LinkedList<Thread> threadQueue) {
        // 由于tryAcquire会多次执行，所以需要判断下当前的thread是否已经在等待队列中
        if (!isQueued(Thread.currentThread())) {
            threadQueue.add(Thread.currentThread());
        }
    }

    @Override
    protected boolean tryRelease(int arg) {
        String localKey = Thread.currentThread().getName();
        // 如果thread队列第一个元素是当前thread，暂定可以激活下一个节点，并将其从该key的threadQueue中移出
        if (lockQueueHolders.get(localKey).getFirst() == Thread.currentThread()) {
            lockQueueHolders.get(localKey).removeFirst();
        } else {
            return false;
        }
        // 获取第一个待激活节点关联的Thread
        Thread firstQueuedThread = getFirstQueuedThread();
        // 如果没有下一个待激活节点，直接返回true
        if (null == firstQueuedThread) {
            return true;
        }
        // 通过getName获取该Thread关联的key
        // 上文中提到ThreadLocal不满足要求，就是在这个地方
        String firstQueuedKey = firstQueuedThread.getName();
        // 如果第一个等待的节点与当前释放锁的节点的key相同，则可以释放
        if (localKey.equals(firstQueuedKey)) {
            return true;
        }
        // 如果下一个待激活的节点关联的Thread，是该key的threadQueue的第一个，表示可以激活
        return firstQueuedThread == lockQueueHolders.get(firstQueuedKey).getFirst();
    }

    @Override
    protected int tryAcquireShared(int arg) {
        return 1;
    }

    @Override
    protected boolean tryReleaseShared(int arg) {
        return true;
    }
}

执行

public class MainRunner {
    public static void main(String[] args) {
        TestCase testCase = new TestCase();
        CaseRunner runner1 = new CaseRunner(testCase, "A", 1, 1000);
        CaseRunner runner2 = new CaseRunner(testCase, "B", 1, 2000);
        CaseRunner runner3 = new CaseRunner(testCase, "C", 1, 3000);
        CaseRunner runner4 = new CaseRunner(testCase, "A", 2, 1000);
        CaseRunner runner5 = new CaseRunner(testCase, "A", 3, 1000);
        CaseRunner runner6 = new CaseRunner(testCase, "C", 2, 3000);
        CaseRunner runner7 = new CaseRunner(testCase, "B", 2, 2000);

        runner1.start();
        runner2.start();
        runner3.start();
        runner4.start();
        runner5.start();
        runner6.start();
        runner7.start();
    }

    static class TestCase {
        MultiReentrantLock multiReentrantLock = new MultiReentrantLock();

        public void tCase(String key, int counter, long sleepTime) throws InterruptedException {
            System.out.println(System.currentTimeMillis() + " key: " + key + '-' + counter + "开始获取锁");
            multiReentrantLock.lock(key);
            System.out.println(System.currentTimeMillis() + " key: " + key + '-' + counter + "获取锁");
            Thread.sleep(sleepTime);
            System.out.println(System.currentTimeMillis() + " key: " + key + '-' + counter + "开始释放锁");
            multiReentrantLock.unlock(key);
            System.out.println(System.currentTimeMillis() + " key: " + key + '-' + counter + "释放锁");
        }
    }

    static class CaseRunner extends Thread {
        private TestCase testCase;
        private String lockKey;
        private long sleepTime;
        private int counter;

        public CaseRunner(TestCase testCase, String lockKey, int counter, long sleepTime) {
            this.testCase = testCase;
            this.lockKey = lockKey;
            this.sleepTime = sleepTime;
            this.counter = counter;
        }

        @Override
        public void run() {
            try {
                testCase.tCase(lockKey, counter, sleepTime);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}

结果打印

1656154306290 key: A-1开始获取锁
1656154306291 key: B-1开始获取锁
1656154306291 key: C-1开始获取锁
1656154306291 key: B-1获取锁
1656154306291 key: A-1获取锁
1656154306291 key: C-1获取锁
1656154306293 key: A-2开始获取锁
1656154306293 key: A-3开始获取锁
1656154306294 key: C-2开始获取锁
1656154306294 key: B-2开始获取锁
1656154307292 key: A-1开始释放锁
1656154307292 key: A-1释放锁
1656154307292 key: A-2获取锁
1656154308292 key: B-1开始释放锁
1656154308292 key: B-1释放锁
1656154308292 key: A-2开始释放锁
1656154308292 key: A-2释放锁
1656154308292 key: A-3获取锁
1656154309293 key: A-3开始释放锁
1656154309293 key: C-1开始释放锁
1656154309293 key: A-3释放锁
1656154309293 key: C-1释放锁
1656154309293 key: C-2获取锁
1656154312299 key: C-2开始释放锁
1656154312299 key: C-2释放锁
1656154312299 key: B-2获取锁
1656154314300 key: B-2开始释放锁
1656154314300 key: B-2释放锁