Demystifying Rust - Hands on implementation of Spin Lock

본 포스트에서는 Spin Lock을 실제로 구현해보고 러스트에서 아래의 개념들을 어떻게 활용하는지를 파헤쳐본다.

• Atomic values
• Lock
• MutexGuard
• UnsafeCell
• Send/Sync

Basic Mechanism for SpinLock

가장 간단한 SpinLock의 구현은 아래코드와 유사한 형태가 될 것이다.

 pub struct SpinLock {
        locked: AtomicBool,
    }

    impl SpinLock {
        pub const fn new() -> Self {
            Self {
                locked: AtomicBool::new(false),
            }
        }
        pub fn lock(&self) {
            while self.locked.swap(true, Ordering::Acquire) {
                std::hint::spin_loop();
            }
        }
        fn unlock(&self) {
            self.locked.store(false, Ordering::Release);
        }
    }

new()의 구현은 차치하고, lock의 구현을 유심히 보자. self.locked.swap()은 swap하기 전의 값을 리턴하고, 첫번째 인자로 값을 교체한다.

Atomic Values

Rust개발자라면 이 지점에서 뭔가 이상한 점을 느낄수도 있겠다.

값에 대한 변경이 있는데 왜 method signature가 &self 인가? &mut self이 되어야 하는거 아닌가?

위의 질문이 해당포스트에서 다루고자 하는 첫번째이다. Atomic value는 그것이 shared reference에서 기인하더라도 변경이 가능하다. 메모리의 실행 순서Memory Ordering을 통해 이뤄낸다.

Lock

그럼 lock 메소드는 정확히 무엇을 수행하는가?
간단히 말하면 말하면 아래와 같다.

"현재 locked 값이 false면 OK", 그게 아니면 무한루프

내부 블록의 std::hint::spin_loop()은 processor에 변경사항이 생길때까지 spinning을 한다는 것을 알림으로써 최적화된 instruction을 유도하는 코드이다. 중요한 부분은 아니니 넘어가도록 하자.

Use basic spin lock in action

위의 아주 베이직한 spin lock을 한번 사용해보자 - 작동하지 않을 것이다.

fn main(){
    let spin_lock = Arc::new(basic::SpinLock::new());

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            spin_lock.lock();
            sleep(Duration::from_secs(1))
        }
    })
    .join();

    spin_lock.lock(); // HANG
}

당연한 결과이다. Background thread에서 lock을 취득한 이후 해당 lock을 해제해주는 unlock 이 호출되지 않았기 때문이다.

아마도 아래와 같은 코드를 짜고 싶은 충동이 들지 모르겠다.

  move || {
            spin_lock.lock();
            sleep(Duration::from_secs(1));
            spin_lock.unlock(); // added
        }

물론, bad practice이다. 우리는 사용자가 우리가 원하는 대로 행동할지 알 수 없다. 사실 unlock의 visibility가 private인 이유도 거기에 있다.

더 좋은 구현은, 아무래도 block scope을 벗어나는 순간 unlock()을 호출하는 방식이 될 것이다:

impl Drop for SpinLock {
     fn drop(&mut self) {
          self.unlock()
      }
  }

위의 코드 자체는 문제가 없다. 그렇지만 아래의 코드는 여전히 Hang이 걸리게 될 것이다.

let spin_lock = Arc::new(basic::SpinLock::new());

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            spin_lock.lock();
            sleep(Duration::from_secs(1));
        }
    })
    .join();

    spin_lock.lock(); // still hang...

What's the problem? - Arc

여러 threads에서 공유할 수 있도록 하기 위해, SpinLock은 Arc에 wrapping된 상태로 각 쓰레드에 전달되었다. 즉, block scope를 벗어날 때 호출되는 Drop은 SpinLock에 대한 것이 아닌 Arc에 대한 것이다. 즉 Reference count를 낮출 뿐, 실제 내부 요소에 접근해서 행위를 하도록 하지 않는다.

Introduction to MutexGuard

자, 아래 조건을 만족하는 구조체를 생각해보자.
1. SpinLock에서 lock을 얻어서 사라지기 전까지만 존재한다.
2. 내부 value로는 SpinLock를 갖는다.

"얻는다"/"사라진다" 라는 말에서 직관적으로 lifetime의 필요성을 느꼈다면 내 기준에서 훌륭한 Rust개발자고 생각이 든다. 아래와 같은 구조체를 생각해볼 수 있겠다.

pub struct MutexGuard<'a> {
    lock: &'a SpinLock,
}

자, 그리고 SpinLock의 lock 메소드 호출시 위의 MutexGuard를 리턴하도록 signature를 변경해보자.

pub fn lock(&self) -> MutexGuard<'_> {
    while self.locked.swap(true, Ordering::Acquire) {
        std::hint::spin_loop();
    }
    MutexGuard { inner: self }
}

자 그리고 나서 Drop trait을 MutexGuard에 적용시켜보자.

impl<'a> Drop for MutexGuard<'a> {
    fn drop(&mut self) {
        self.inner.unlock()
    }
}

그리고 아래처럼 실행시켜보자.

fn main() {
    let spin_lock = Arc::new(basic::SpinLock::new());

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            spin_lock.lock();
            sleep(Duration::from_secs(1));
        }
    })
    .join();

    spin_lock.lock(); // not hang

    println!("Success!") 

SpinLock with value

위의 코드는 SpinLock의 동작 매커니즘을 돕지만, 아래 두가지에 대한 대한기술이 부족하다.

• 내부 데이터에 대한 접근을 어떻게 이뤄내고
• 해당 데이터를 어떻게 변경 할 수 있는지

하지만 그 전에, 아래 질문에 대한 답을 먼저 해보자.

현재 구현되어있는 lock과 unlock, 그리고 MutexGaurd로, 우리가 어떤 조작을 하던 Spinlock에 대한 접근은 순차적임을 보장할 수 있는가?

정답은 Yes이다. 그리고 그 답변에 확신이 있다면, unsafe 또한 현재 구현안에서 안전하다는 것을 확신한다는 것과 같다.

우리가 원하는 것은 아래와 같다.

• 불특정 타입 T를, SpinLock이 내부 value로 가질 수 있도록 하고
• 값에 대한 접근이 가능하며
• 변경또한 가능하고
• 각기 다른 thread에서 접근하더라도 Thread safety를 보장

Change to Structs

불특정 타입 T를 받도록 하기 위해, Generic을 사용한다.

pub struct MutexGuard<'a, T> {
    inner: &'a SpinLock<T>,
}

pub struct SpinLock<T> {
    locked: AtomicBool,
    inner: T,
}

Change to impl

impl<T> SpinLock<T> {
    pub const fn new(val: T) -> Self {
        Self {
            locked: AtomicBool::new(false),
            inner: val,
        }
    }
    pub fn lock(&self) -> MutexGuard<'_, T> {
        while self.locked.swap(true, Ordering::Acquire) {
            std::hint::spin_loop();
        }
        MutexGuard { inner: self }
    }
    pub fn unlock(&self) {
        self.locked.store(false, Ordering::Release);
    }
}

impl<'a, T> Drop for MutexGuard<'a, T> {
    fn drop(&mut self) {
        self.inner.unlock()
    }
}

implement Deref for better ergonomics

inner value들 (MutexGuard의 inner혹은 SpinLock의 inner)의 visibility가 private인것은 사뭇 직관적이다.

그러나 그러한 구현은, MutexGuard를 통해서 값에 대한 참조를 하는 것을 원천적으로 차단하고, 심지어 pub으로 변환한다고 해도, guard.inner.inner의 방식을 강제하기에 매우 사용성이 떨어진다.

따라서 Deref를 통해, Guard를 사용하는 것이지만 아래와 같이 마치 실제 SpinLock의 value에 직접 접근하는 듯한 경험을 주고자 한다:

fn main() {
    let spin_lock = Arc::new(basic::SpinLock::new(1));

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            let guard = spin_lock.lock();
            sleep(Duration::from_secs(1));
            println!("{}", *guard) // like this!
        }
    })
    .join();

    spin_lock.lock(); // not hang

    println!("Success!")
}

Deref 실제 구현은 아래와 같다

impl<'a, T> Deref for MutexGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.inner.inner
    }
}

Modifying value and Send/Sync

아래와 같이 값의 변경과 Thread간 순서보장을 가능하게 하고 싶다고 가정해보자.

fn main() {
    let spin_lock = Arc::new(basic::SpinLock::new(1));

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            let mut guard = spin_lock.lock();
            *guard = 5; // change to value
            sleep(Duration::from_secs(1));
        }
    })
    .join();

    let guard = spin_lock.lock(); 

    assert_eq!(*guard, 5); // test!
}

우리가 시도해볼 수 있는 첫번째는, 일단 DerefMut을 짜보는 것이다.

impl<'a, T> DerefMut for MutexGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner.inner
    }
}

위의 코드는 컴파일하지 않는다. 이유는 간단하다. MutexGuard가 inner를 shared reference의 형태로 받기 때문이다.

그리고 이 시점에서 위에서 던진 질문에 대한 답을 다시 한번 상기시켜보자.

현재 구현되어있는 lock과 unlock, 그리고 MutexGaurd로, 우리가 어떤 조작을 하던 Spinlock에 대한 접근은 순차적임을 보장할 수 있는가?

Yes라는 것에 확신을 얻었다면, 우리는 shared reference를 통해서도 내부 값을 변경할 수 있는 메커니즘을 생각해봐야한다. 그리고 Rust에서, Low레벨 단계에서 해당 매커니즘을 제공하는 API는 UnsafeCell이 유일하다.

즉, SpinLock이 내부 value를 갖는 형태는 단순히 T가 아니라 UnsafeCell<T>가 되어야 한다는 의미이다.

Changes in struct and impl

    pub struct SpinLock<T> {
        locked: AtomicBool,
        inner: UnsafeCell<T>, // change!
    }

    impl<T> SpinLock<T> {
        pub const fn new(val: T) -> Self {
            Self {
                locked: AtomicBool::new(false),
                inner: UnsafeCell::new(val), // change!
            }
        }
    ...
    
    }

Change in Deref and DerefMut

UnsafeCell로 wrapping된 value에 대한 접근은 아래와 같은 단계로 이뤄진다.

• unsafe 블록을 지정하여 해당 코드가 unsafe하다는 것 을 알리고, 컴파일러로 하여금 해당 라인에 대한 safety check을 하지 않도록 한다.
• get() method를 통해 pointer를 얻어낸다.
• dereference operator (*) 를 사용, 실제 값을 얻어낸다.
• shared reference가 필요한 경우 &을 붙여주고 mutable reference가 필요다면 &mut 을 붙여준다.

위의 단계를 따르면, Deref와 DerefMut는 아래와 같이 구현될 것이다.

    impl<'a, T> Deref for MutexGuard<'a, T> {
        type Target = T;
        fn deref(&self) -> &Self::Target {
            unsafe { &*self.inner.inner.get() } 
        }
    }

    impl<'a, T> DerefMut for MutexGuard<'a, T> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            unsafe { &mut *self.inner.inner.get() }
        }
    }

자, 이제 대충 준비가 끝난 것으로 보이지..만 아래의 코드는 동작하지 않는다.

fn main() {
    let spin_lock = Arc::new(basic::SpinLock::new(1));

    let j = spawn({ // Error! 
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            let mut guard = spin_lock.lock();
            sleep(Duration::from_secs(1));
            *guard = 5;
        }
    })
    .join();

    let guard = spin_lock.lock(); // not hang

    assert_eq!(*guard, 5);
}

unsafe Send/Sync

그것은 SpinLock에 주어지는 Generic type T에 대해 thread간 share가 가능한지, send가 가능한지를 명시해주지 않았기 때문이다. 아래의 코드 한줄로 해당 문제는 해결된다.

unsafe impl<T> Sync for SpinLock<T> where T: Send {}

Sync와 Send traits에 대한 impl은 그 자체로 unsafe하다. 따라서 해당 impl 전에 unsafe마킹을 해주어야 한다.

Conclusion

아래는 전체 코드 예제이다.

pub struct SpinLock<T> {
    locked: AtomicBool,
    inner: UnsafeCell<T>,
}

impl<T> SpinLock<T> {
    pub const fn new(arg: T) -> Self {
        Self {
            locked: AtomicBool::new(false),
            inner: UnsafeCell::new(arg),
        }
    }
    pub fn lock(&self) -> MutexGuard<T> {
        while self.locked.swap(true, Ordering::Acquire) {
            std::hint::spin_loop();
        }
        MutexGuard { lock: self }
    }
    fn unlock(&self) {
        self.locked.store(false, Ordering::Release);
    }
}

pub struct MutexGuard<'a, T> {
    lock: &'a SpinLock<T>,
}

impl<T> Deref for MutexGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { &*self.lock.inner.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.lock.inner.get() }
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        self.lock.unlock();
    }
}

unsafe impl<T> Sync for SpinLock<T> where T: Send {}

fn main() {
    let spin_lock = Arc::new(SpinLock::new(1));

    let j = spawn({
        let spin_lock = Arc::clone(&spin_lock);
        move || {
            let mut guard = spin_lock.lock();
            sleep(Duration::from_secs(1));
            *guard = 5;
        }
    })
    .join();

    let guard = spin_lock.lock(); // not hang

    assert_eq!(*guard, 5);
}