【练习/JUC】实现ReentrantLock简单版

实现了一个精简版的ReentrantLock，逻辑比较清晰。思路应该和真正的ReentrantLock差不多？希望发现Bug的读者指点下。

public class Lock {
    volatile Thread owner;
    volatile Node head = new Node();
    AtomicReference<Node> tail = new AtomicReference<>(head);
    AtomicInteger state = new AtomicInteger(0);
    AtomicInteger queueLength = new AtomicInteger(0);

    /**
     * head(?) -> Node(t1) -> Node(t2)
     */
    static class Node {
        volatile Node next;
        volatile Thread thread;

        Node() {
        }

        Node(Node next, Thread thread) {
            this.next = next;
            this.thread = thread;
        }
    }

    public void lock() {
        Thread t = Thread.currentThread();
        Node cur = new Node(null, t);
        while (!state.compareAndSet(0, 1)) { // 循环获取锁
            // 锁重入
            if (owner == t) {
                state.incrementAndGet();
                return;
            }
            // 获取锁失败
            Node prev = null;
            Node next = head.next;
            if (next == null || next.thread != t) {
                // 入队
                // 这里即使多个线程一起入队也是安全的
                // 因为CAS更换tail成功才会更新记录的前一个节点next
                // 即使next还未更新，其他线程入队CAS已经是新的tail
                // 之后依次拉链就可以了
                do {
                    prev = tail.get();
                } while (!tail.compareAndSet(prev, cur));
                prev.next = cur;
                // 额外记录了一下队列长度，可以忽略
                queueLength.incrementAndGet();
            }

            // 挂起当前线程
            // 如果是队列首部再次尝试获取
            // 主要防止入队过程中，还没有给head设置好next
            // 持有锁的线程释放了锁无法唤醒队列头线程
            // 在这里已经拉链好了，再尝试获取一次就可以解决这个问题
            if (prev == head && state.get() != 0)
                LockSupport.park();
        }
        // 这之后是获取锁的单线程执行
        // 在队列中就出队
        if (head.next != null && head.next.thread == t) {
            head = head.next;
            queueLength.decrementAndGet();
        }
        owner = t;
    }

    public void unlock() {
        Thread t = Thread.currentThread();
        // 没持有锁
        if (owner != t) {
            throw new RuntimeException(t.getName() + " not have such lock");
        }

        int state = this.state.get();
        // 锁全部退出
        if (state == 1) {
        	// 必须先设置owner为null再释放(state--)
        	// 防止释放后阻塞，其他线程获取到锁修改owner
        	// 之后恢复把owner改成了null导致获取锁的线程释放时进入上面的没持有锁异常
            owner = null;
            // 唤醒头节点线程
            this.state.decrementAndGet();
            Node next = head.next;
            if (next != null) {
                LockSupport.unpark(next.thread);
            }
        }

    }
}

测试代码：

class Test {
    static int total = 10000;

    public static void main(String[] args) throws InterruptedException {
        Lock lock = new Lock();
        Runnable runnable = () -> {
            while (total > 0) {
                try {
                    Thread.sleep(10);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
                lock.lock();
                try {
                    if (total > 0) {
                        System.out.println(Thread.currentThread().getName() + ": " + --total);
                    }
                } finally {
                    lock.unlock();
                }
            }
        };
        List<Thread> list = new ArrayList<>();
        for (int i = 0; i < 20; i++) {
            Thread thread = new Thread(runnable);
            list.add(thread);
            thread.setDaemon(true);
            thread.start();
        }
        list.forEach(t -> {
            try {
                t.join();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        });
    }
}

经过测试没什么问题。