1. 程式人生 > 其它 >自己手動寫一個鎖Lock

自己手動寫一個鎖Lock

準備

自己實現一個鎖我們要準備一個變數、一個佇列、Unsafe類來執行CAS/park/unpark的Unsafe類。

一個變數

​ 一個變數的值為1的時候就說明已加鎖,變數值為0的時候就說明未加鎖

一個佇列

​ 多個執行緒對同一個鎖的爭用肯定只有一個能成功,沒有搶到鎖的其它的執行緒就要排隊,所以我們還需要一個佇列

Unfafe類

​ 多個執行緒只能有一個執行緒把變數的值修改為1,且當它的值為1的時候其它執行緒不能再修改它的值,我們使用Unsafe這個類來做CAS操作。而且排隊的執行緒不能繼續往下執行,只能阻塞。我們用Unsfae類來阻塞(park)和喚醒(unpark)執行緒。

大體的流程示意圖如下:

詳細程式碼

主要程式碼

public class MyLock {

    //用來標記當前鎖是否被執行緒佔用,0:不佔用 1:不佔用
    private volatile int state;

    private Node empty = new Node();

    // 佇列頭
    private volatile Node head;
    // 佇列尾
    private volatile Node tail;

    public MyLock() {
        head = tail = empty;
    }

    private static Unsafe unsafe;

    // state的偏移量
    private static long stateOffset;

    // tail的偏移量
    private static long tailOffset;

    static {
        try {
            // 獲取Unsafe的例項
            Field f = Unsafe.class.getDeclaredField("theUnsafe");
            f.setAccessible(true);
            unsafe = (Unsafe) f.get(null);
            // 獲取state的偏移量
            stateOffset = unsafe.objectFieldOffset(MyLock.class.getDeclaredField("state"));
            // 獲取tail的偏移量
            tailOffset = unsafe.objectFieldOffset(MyLock.class.getDeclaredField("tail"));
        } catch (NoSuchFieldException e) {
            e.printStackTrace();
        } catch (IllegalAccessException e) {
            e.printStackTrace();
        }
    }

    // 原子更新state欄位
    private boolean compareAndSetState(int expect, int update) {
        return unsafe.compareAndSwapInt(this, stateOffset, expect, update);
    }

    // 原子更新tail欄位
    private boolean compareAndSetTail(Node expect, Node update) {
        return unsafe.compareAndSwapObject(this, tailOffset, expect, update);
    }
  
    // 加鎖
    public void lock() {
        // 嘗試更新state欄位,更新成功說明佔有了鎖
        if (compareAndSetState(0, 1)) {
            return;
        }
        // 未更新成功入隊
        Node node = enq();
        Node prev = node.prev;
        // 入隊之後,再次嘗試獲取鎖
        // 需要檢測當前入隊節點的上一個節點是不是head,不是則阻塞。
        while (node.prev != head || !compareAndSetState(0, 1)) {
            // 未獲取到鎖,阻塞
            unsafe.park(false, 0L);
        }
        // 下面不需要原子更新,因為同時只有一個執行緒訪問到這裡
        // 獲取到鎖了且上一個節點是head
        // head後移一位
        head = node;
        // 清空當前節點的內容,協助GC回收
        node.thread = null;
        // 將上一個節點從連結串列中剔除,協助GC回收
        node.prev = null;
        prev.next = null;
    }

    // 解鎖
    public void unLock() {
        // 把state更新成0,這裡不需要原子更新,因為同時只有一個執行緒訪問到這裡
        state = 0;
        // 下一個待喚醒的節點
        Node next = head.next;
        // 下一個節點不為空,就喚醒它
        if (next != null) {
            unsafe.unpark(next.thread);
        }
    }

    /**
     * 沒有搶到鎖的執行緒入佇列(入隊過程是原子的)
     */
    private Node enq() {
        while (true) {
            Node t = tail;
            Node node = new Node(Thread.currentThread(), t);
            if (compareAndSetTail(t, node)) {
                node.prev = t;
                t.next = node;
                return node;
            }
        }
    }
}

測試類

public class TestMyLock {

    private static int count = 0;

    @Test
    public void testMyLock() throws InterruptedException {
        MyLock lock = new MyLock();
        CountDownLatch countDownLatch = new CountDownLatch(1000);

        IntStream.range(0, 1000).forEach(i -> new Thread(() -> {
            lock.lock();
            try {
                IntStream.range(0, 10000).forEach(j -> {
                    count++;
                });
            } finally {
                lock.unLock();
            }
            countDownLatch.countDown();
        }, "tt" + i).start());

        countDownLatch.await();
        System.out.println(count);
    }
}

測試結果

未加鎖列印:

測試加鎖後,是我們要的結果,鎖有效

詳細呼叫過程

假設現在有3個執行緒進入lock方法

​ 執行緒1原子更新state成功然後去執行下面的邏輯。執行緒2和執行緒3沒有搶到鎖,都阻塞在unsafe.park(false, 0L)處。並且此時佇列是這樣的:Node(empty)—Node(執行緒2)—Node(執行緒3)

​ 執行緒1執行業務程式碼結束並執行unLock方法,執行unLock方法後,在unsafe.park(false, 0L)處阻塞的執行緒

​ 執行緒2線上程1執行unLock方法的unsafe.unpark(next.thread)被喚醒,再執行一次while (node.prev != head || !compareAndSetState(0, 1)),發現比對符合條件,執行下面的程式碼。並且佇列現在變成Node(執行緒2)—Node(執行緒3)

接下來就是重複上面的呼叫過程。