使用AtomicReference类通过自旋和CAS简单实现自定义可重入锁

AtomicReference是原子引用并发类，内部通过调用Unsafe的本地方法CAS实现。下面我们使用CAS和自旋实现一个简单的可重入锁。

自旋省去了线程的切换，但是空循环增加了对CPU的消耗。

首先，定义MyReentrantLock类，有lock加锁和unlock释放锁，加锁设置超时时间，AtomicReference内部保存当前获取到锁线程的ID。

import java.util.concurrent.atomic.AtomicReference;

public class MyReentrantLock {

    private AtomicReference<Long> ar = new AtomicReference<>();
    private int lockCount;

    public boolean lock(int timeoutMillisecond) {
        long startTime = System.currentTimeMillis();
        boolean lockResult = false;
        for (; ; ) {
            if (ar.get() != null && Thread.currentThread().getId() == ar.get()) { // 线程重入
                lockCount++;
                lockResult = true;
            }
            if (ar.compareAndSet(null, Thread.currentThread().getId())) { // 获取锁
                lockCount = 1;
                lockResult = true;
            }

            if (System.currentTimeMillis() - startTime > timeoutMillisecond) { // 超时
                System.out.println(Thread.currentThread().getId() + "线程获取锁超时");
                return false;
            }
            if (lockResult) {
                System.out.println(Thread.currentThread().getId() + "线程获取锁成功: " + lockCount);
                return lockResult;
            }
        }
    }

    public void unlock() {
        if (ar.get() != null) {
            if (Thread.currentThread().getId() != ar.get()) { // 线程重入
                throw new RuntimeException("Thread not got lock");
            }
            if (lockCount > 0) {
                lockCount--;
                System.out.println(Thread.currentThread().getId() + "线程释放锁成功: " + lockCount);
            }
            if (lockCount == 0) {
                ar.compareAndSet(Thread.currentThread().getId(), null);
            }
        }
    }

}

测试类

public class MyLockTest {

    private static MyReentrantLock myLock = new MyReentrantLock();

    public static void main(String[] args) {
        // 简单使用
        new Thread(() -> {
            if (myLock.lock(1000)) {
                try {
                    Thread.sleep(1500); // 模拟线程操作时间
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                myLock.unlock();
            }
        }).start();

        // 递归重入锁
        new Thread(() -> {
            reentrant(1);
        }).start();

    }

    private static void reentrant(int n) {
        if (n > 5) {
            return;
        }
        if (myLock.lock(1000)) {
            try {
                reentrant(++n);
                Thread.sleep(100);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            myLock.unlock();
        }
    }

}

执行结果

Thread.sleep(1500)运行结果
14线程获取锁成功: 1
15线程获取锁超时
14线程释放锁成功: 0


Thread.sleep(100)运行结果
14线程获取锁成功: 1
14线程释放锁成功: 0
15线程获取锁成功: 1
15线程获取锁成功: 2
15线程获取锁成功: 3
15线程获取锁成功: 4
15线程获取锁成功: 5
15线程释放锁成功: 4
15线程释放锁成功: 3
15线程释放锁成功: 2
15线程释放锁成功: 1
15线程释放锁成功: 0