AQS 源码探究_01 手写一个简化的 ReentrantLock 可重入锁

完整代码

Lock

/**
 * @author wcc
 * @date 2022/2/10 22:13
 */
public interface Lock {

  void lock();

  void unlock();

}

MinReentrantLock

import sun.misc.Unsafe;

import java.lang.reflect.Field;
import java.util.concurrent.locks.LockSupport;

/**
 * @author wcc
 * @date 2022/2/10 22:13
 */
public class MinReentrantLock implements Lock{

  /**
   * 锁的是什么？
   * 资源 -> state
   * 0：表示未加锁状态
   * >0：表示当前 lock 是加锁状态
   */
  private volatile int state;

  /**
   * 独占模式：
   * 同一时刻只有一个线程可以持有锁，其他的线程在未获取到锁的时候，会被阻塞
   */

  // 表示当前独占锁的线程（占用锁的线程）
  private Thread exclusiveOwnerThread;

  /**
   * 需要有两个引用去维护阻塞队列
   * 1.head 指向队列的头结点
   * 2.tail 指向队列的尾节点
   */
  // 比较特殊：head节点对应的线程就是当前占用锁的线程
  private Node head;
  private Node tail;

  /**
   * 阻塞的线程被封装成什么？
   * Node 节点，然后放入到FIFO 队列，先进先出队列
   */
  static final class Node{

    // 前置节点引用
    Node prev;

    // 后置节点引用
    Node next;

    // 封装的线程本身
    Thread thread;

    public Node(Thread thread) {
      this.thread = thread;
    }

    public Node(){

    }
  }

  /**
   * 获取锁
   * 假设当前锁被占用，则会阻塞调用者线程直到抢占到锁为止
   * 模拟公平模式的公平锁
   * 什么是公平锁？讲究一个先来后到
   *
   * lock 的过程是怎么样的？
   * 情景一：线程进来后发现，当前state == 0，这个时候就很幸运，直接抢锁
   * 情景二：线程进来后发现，当前state > 0，这个时候就需要将当前线程入队
   */
  @Override
  public void lock() {
    // 第一次抢占锁成功：将state 设置为1
    // 第n次重入锁的时候将state设置为n
    acquire(1);
  }

  public int getState() {
    return state;
  }

  public Thread getExclusiveOwnerThread() {
    return exclusiveOwnerThread;
  }

  public Node getHead() {
    return head;
  }

  public Node getTail() {
    return tail;
  }

  /**
   * 竞争资源
   * 1.尝试获取锁，成功则占用锁且返回
   * 2.抢占锁失败，阻塞当前线程
   * @param arg
   */
  private void acquire(int arg){
    if(!tryAcquire(arg)){
      // 更复杂的逻辑了 1.需要将当前线程封装成node节点加入到阻塞队列中
      Node node = addWaiter();
      // 2.需要将当前线程给park掉，使得线程处于挂起状态
      acquireQueued(node, arg);
    }
  }

  /**
   * 尝试抢占锁失败需要做什么呢？
   * 1.需要将当前线程封装成node节点加入到阻塞队列中
   * 2.需要将当前线程给park掉，使得线程处于挂起状态
   *
   * 唤醒后呢？
   * 1.需要检查当前node节点是否为head.next节点
   *  (head.next 节点是拥有抢占权限的线程，其他node节点线程都没有抢占权限)
   * 2.抢占锁
   *  成功：1.将当前node设置为head，将老的head做出队操作，返回到业务层面
   *  失败：2.继续park，等待被唤醒
   *
   * ==>
   * 1.添加到阻塞队列的逻辑，addWaiter()
   * 2.在队列中竞争资源的逻辑  acquireQueued()
   */

  /**
   * 当前线程节点入队的方法
   * @return 返回当前线程对应的node节点
   *
   * addWaiter()方法执行完毕之后，保证当前线程已经入队成功
   */
  private Node addWaiter(){
    Node newNode = new Node(Thread.currentThread());
    // 如何入队呢？
    // 1.找到 newNode的前置节点
    // 2.更新newNode的prev = 前置节点 next = null
    // 3.CAS 更新tail为newNode
    // 4.更新 pred.next = newNode

    Node pred = tail;

    // 前置条件：队列已经有等待者node了，当前newNode不是第一个入队的节点
    if(pred != null){
      newNode.prev = pred;
      // 条件成立：说明当前线程成功入队
      if(compareAndSetTail(pred, newNode)){
        pred.next = newNode;
        return newNode;
      }
    }

    // 执行到这里，有几种情况：应该执行入队操作
    // 1.tail == null：说明当前队列是空队列
    // 2.CAS 设置当前newNode 为 tail的时候失败了，被其他线程抢先一步了
    enq(newNode);
    return newNode;
  }

  /**
   * 需要将当前线程给park掉，使得线程处于挂起状态
   *
   * 唤醒后呢？
   * 1.需要检查当前node节点是否为head.next节点
   *  (head.next 节点是拥有抢占权限的线程，其他node节点线程都没有抢占权限)
   * 2.抢占锁
   *  成功：1.将当前node设置为head，将老的head做出队操作，返回到业务层面
   *  失败：2.继续park，等待被唤醒
   */
  private void acquireQueued(Node node, int arg){
    // 只有当前node线程成功获取到锁以后才会跳出自旋操作
    retry:
    for (;;){
      // 什么情况下，当前node被唤醒之后可以尝试去获取锁呢？
      // 只有一种情况，当前node是head的后继节点，才有这个权限
      // 不就是先来后到吗？

      // head 节点就是当前持锁节点
      Node pred = node.prev;
      // 条件一：pred == head 成立：说明当前node拥有抢占权限
      if(pred == head && tryAcquire(arg)){
        // 这里面，说明当前线程竞争锁成功了
        // 需要做点什么？
        // 1.设置当前head为当前线程的node节点
        // 2.协助原始 head 出队操作
        setHead(node);

        pred.next = null; // help GC
        return;
      }

      System.out.println("线程：" + Thread.currentThread().getName() + "park...");
      // park.. 让当前线程处于挂起状态
      LockSupport.park();

      // 什么唤醒被park的线程呢？
      // unlock 过程了!
      System.out.println("线程：" + Thread.currentThread().getName() + "unpark...");
    }

  }

  /**
   * 自旋入队，只有成功后才返回
   * 1.tail == null：说明当前队列是空队列
   * 2.CAS 设置当前newNode 为 tail的时候失败了，被其他线程抢先一步了
   * @param node
   * @return
   */
  private void enq(Node node){
    retry:
    for (;;){
      // 第一种情况：队列是空
      // ==> 说明当前线程是第一个抢占锁失败的线程
      // 当前持有锁的线程并没有设置过任何node，所以作为该线程的第一个后驱节点需要给它擦屁股
      // 1.给当前持有锁的线程补充一个 node节点作为head节点入队，head节点任何时候都代表当前占用锁的线程
      if(tail == null){
        // 条件成立：说明当前线程给当前持有锁的线程补充了head头结点操作成功
        if(compareAndSetHead(new Node())){
          tail = head;
          // 注意，这里并没有直接返回，而是继续进行自旋操作
        }
        // 到这里的情况只有：CAS 设置当前newNode 为 tail的时候失败了，被其他线程抢先一步了
      }else{
        // 说明当前队列中已经有head了，这里是一个追加node节点的过程
        // 如何入队呢？
        // 1.找到 newNode的前置节点
        // 2.更新newNode的prev = 前置节点 next = null
        // 3.CAS 更新tail为newNode
        // 4.更新 pred.next = newNode

        Node pred = tail;

        // 前置条件：队列已经有等待者node了，当前newNode不是第一个入队的节点
        if(pred != null){
          node.prev = pred;
          // 条件成立：说明当前线程成功入队
          if(compareAndSetTail(pred, node)){
            pred.next = node;
            // 注意，入队成功之后一定要return
            return;
          }
        }
      }
    }
  }

  /**
   * 尝试获取锁，不会阻塞线程
   * true：抢占成功
   * false：抢占失败
   */
  private boolean tryAcquire(int arg){
    if(state == 0){
      // 当前 state == 0的时候，是否可以直接抢占锁呢？
      // 不可以，因为咱们模拟的是公平锁，先来后到...
      // 条件一：!hasQueuePredecessor() 取反之后返回 true：表示当前线程前面没有等待者线程
      // 条件二：compareAndSetState(0,arg) 为什么使用CAS 呢？
      //  因为 lock 方法可能有多线程调用的情况，需要使用CAS方式
      // 成立：说明当前线程抢占锁成功
      if(!hasQueuePredecessor() && compareAndSetState(0,arg)){
        // 抢占锁成功，需要干点什么
        // 1.需要将 exclusiveOwnerThread 设置为当前进入if块中的线程
        this.exclusiveOwnerThread = Thread.currentThread();
        return true;
      }

    // 什么时候会执行 else-if 条件？
    // 当前锁是被占用的时候会来到这个条件这里..
    // 条件成立：Thread.currentThread() == this.exclusiveOwnerThread
    // 说明当前线程即为持有锁的线程，是需要返回true的，可以继续占用锁，可重入锁，并且更新state值
    }else if(Thread.currentThread() == this.exclusiveOwnerThread){

      // 为什么这里不使用CAS的方式来更新state的值？
      // 因为这里并没有并发，只有一个持锁线程能够进入

      // 可重入锁的流程：
      // 说明当前线程即为持有锁的线程，是需要返回true的
      int currentState = getState();
      currentState = currentState + arg;
      // 越界判断..，这里为了简便，就不再做了
      this.state = currentState;
      return true;
    }
    // 什么时候会返回false？
    // 1.CAS 加锁失败或者当前线程不是队列中的头部线程
    // 2.state > 0，且当前线程不是持锁线程
    return false;
  }

  /**
   * true：表示当前线程在队列中有等待者线程
   * flase：表示当前线程前面没有其他等待者线程
   *
   * 调用链：lock -> acquire -> tryAcquire -> (state == 0) hasQueuePredecessor
   *
   * 什么时候返回false呢？
   * 1.当前队列是空
   * 2.当前线程是head.next节点线程  head.next节点线程在任何时候都有权力去争取lock
   */
  private boolean hasQueuePredecessor(){
    Node h = head;
    Node t = tail;
    Node s;

    // 条件一：h != t 成立：说明当前队列已经有元素了
    //  不成立：1. h == t == null
    //         2.h == t == head 第一个获取锁失败的线程会为当前持有锁的线程补充创建一个head节点

    // ((s = h.next) == null || s.thread != Thread.currentThread()) 前置条件：条件一成立
    // 排除几种情况：
    // 条件2.1：(s = h.next) == null
    //  极端情况：第一个获取锁失败的线程会为持锁的线程补充创建head头结点，然后再自旋入队 1.casTail 2.pred(node).next = node
    //  其实想表达的就是：已经有head.next节点了，其他线程在来这的时候需要返回true
    // 条件2.2：s.thread != Thread.currentThread() 前置条件：head.next 不为空
    //  条件成立：说明当前线程就不是head.next节点对应的线程，返回true
    //  条件不成立：说明当前线程就是head.next 节点对应的线程，需要返回false，回头线程会去获取锁了
    return h != t && ((s = h.next) == null || s.thread != Thread.currentThread());
  }

  /**
   * 释放锁
   */
  @Override
  public void unlock() {
    release(-1);
  }

  private void release(int arg){
    // 条件成立：说明线程已经完全释放锁了
    // 需要干点啥呢？
    // 阻塞队列里面还有好几个睡觉的线程呢？是不是应该唤醒一个线程呢？
    if(tryRelease(arg)){
      Node head = this.head;

      // 你需要知道有没有等待者？就是判断head.next == null 说明没有等待者，head.next != null 说明有等待者
      if(head.next != null){
        // 公平锁，就是唤醒head.next 节点
        unparkSuccessor(head);
      }
    }
  }

  /**
   * 完全释放锁成功则返回true，否则说明当前state仍是>0的，返回false
   * @return
   */
  private boolean tryRelease(int arg) {
    int c = getState() - arg;

    if(getExclusiveOwnerThread() != Thread.currentThread()){
      throw new RuntimeException("fuck you! must getLock");
    }

    // 如果执行到这里，存在并发吗？ 不存在，只有一个线程 exclusiveOwnerThread 会来到这里

    // 条件成立：说明当前线程持有的lock锁已经完全释放了
    if(c == 0){
      // 这个时候需要做什么呢？
      // 1.exclusiveOwnerThread = null
      // 2.设置state为0
      this.exclusiveOwnerThread = null;
      this.state = c;
      return true;
    }

    this.state = c;

    return false;
  }

  private void unparkSuccessor(Node node){
    Node s = node.next;
    if(s != null && s.thread != null){
      LockSupport.unpark(s.thread);
    }
  }

  private void setHead(Node node){
    this.head = node;
    // 为什么？因为当前node已经是获取锁成功的线程了...
    node.thread = null;
    node.prev = null;
  }


  /**
   * 基于CAS 的方式，得到state、head、tail 属性的setter 方法
   */
  private static Unsafe unsafe;
  private static long stateOffset;
  private static long headOffset;
  private static long tailOffset;

  static {
    try {
      Field f = Unsafe.class.getDeclaredField("theUnsafe");
      f.setAccessible(true);
      unsafe = (Unsafe) f.get(null);

      stateOffset = unsafe.objectFieldOffset
          (MinReentrantLock.class.getDeclaredField("state"));
      headOffset = unsafe.objectFieldOffset
          (MinReentrantLock.class.getDeclaredField("head"));
      tailOffset = unsafe.objectFieldOffset
          (MinReentrantLock.class.getDeclaredField("tail"));
    }catch (Exception e){
      e.printStackTrace();
    }
  }

  private final boolean compareAndSetHead(Node update){
    return unsafe.compareAndSwapObject(this, headOffset, null, update);
  }

  private final boolean compareAndSetTail(Node expect, Node update){
    return unsafe.compareAndSwapObject(this, tailOffset, expect, update);
  }

  private final boolean compareAndSetState(int expect, int update){
    return unsafe.compareAndSwapObject(this, stateOffset, expect, update);
  }
}