바른 `WaitGroup` 사용법

간혹 여러 고루틴(goroutine)을 동시에 실행시키고, 이 고루틴들이 모두 종료될 때까지 기다리고 싶을 때가 있다. 고루틴은 단순히 주어진 작업을 진행하고 끝나면 종료되는 하나의 루틴일 뿐이기에, 이러한 추가 기능이 필요하면 따로 구현해야 한다.

WaitGroup은 바로 이런 기능을 아주 쉽게 구현할 수 있도록 도와준다. WaitGroup의 내부에선 기다릴 특정 고루틴의 개수를 저장하고 있다. 기다리고 싶은 고루틴을 생성할 때마다 WaitGroup의 Add메서드를 호출해서 기다릴 고루틴의 수를 늘리고, 반대로 고루틴이 종료되면 Done메서드를 호출해서 그 수를 줄이는 방식으로 사용한다. 메인 고루틴에서는 Wait메서드를 호출해 모든 고루틴이 종료되어 내부의 대기 고루틴의 수가 0이 될 때까지 기다리게 할 수 있다.

쉽고 간단한 방법이지만 병행 프로그래밍(concurrent programming)에서 실수하기 쉬워 주의가 필요하다.

var wg sync.WaitGroup
for i := 0; i < 10; i++ {
    go func() {
        wg.Add(1) // 대기 고루틴 수 1 증가
        // 특정 작업을 진행...
        wg.Done() // 대기 고루틴 수 1 감소
    }()
}
wg.Wait() // 자식 고루틴 10개가 모두 종료될 때까지 대기

딱히 문제 될 게 없어 보인다. 고루틴을 생성할 때마다 WaitGroup의 고루틴 대기 개수를 증가시키고 있고 작업이 완료되면 알아서 개수를 감소시키도록 하고 있다. 따라서 마지막의 Wait메서드는 10개의 고루틴이 모두 작업을 완료될 때까지 기다리다가(blocking) 완료가 확인되면 그다음 작업을 진행할 것이다.

하지만, 실제로는 그렇지 않을 때도 있다. 사실 고루틴이 어떻게 스케쥴링될지는 실행할 때마다 다르기 때문에, 생성된 고루틴들의 스케쥴링이 늦어질 경우 예상치 못한 결과를 얻을 수 있다. 무슨 말인가 하면 예를 들어 위에서 고루틴이 10개 생성되고 고루틴 모두가 작업을 시작하기 전에, 즉 모든 고루틴의 wg.Add(1)가 호출되기 전에 메인 고루틴이 마지막 라인(Wait호출 부분)에 도달할 수 있다는 것이다. 아직 대기 개수가 증가(Add)되지 않아 초기값인 0인 그대로이고 메인 고루틴은 '대기할 고루틴이 없구나'하고 자식 고루틴들을 기다리지 않고 바로 다음 작업을 진행해버리게 된다.

이 문제를 해결하기 위해선, 언제 스케쥴링돼서 실행될지 모를 고루틴안에서 대기 개수를 증가시키는 게 아니라 해당 고루틴을 실행하기 직전에 증가시켜야 한다.

var wg sync.WaitGroup
for i := 0; i < 10; i++ {
    wg.Add(1) // 대기 고루틴 수 1 증가
    go func() {
        defer wg.Done() // 대기 고루틴 수 1 감소

        // 특정 작업을 진행...
    }()
}
wg.Wait() // 자식 고루틴 10개가 완료될 때까지 대기

여기서 한 가지 더 눈여겨볼 점은 완료 시에 대기 고루틴 수를 감소시키는 코드로 defer를 이용했다는 것이다. 이렇게 하면 고루틴이 마지막에 도달하기 전에 다른 조건으로 도중에 종료되더라도 대기 고루틴 수를 감소시킬 수 있다.

바른 WaitGroup 사용법

바른 `WaitGroup` 사용법