How to use sync.Once for executing a function exactly once

Photo by Obi Onyeador on Unsplash

I was working on a project and I wanted something to do the these things for me:

1. Execute a bunch of statements (a function).
2. from multiple go routines.
3. only once.
4. and at least once before the parent go routine exits.

Consider a simple function like below, which can be called from multiple locations, but the code has to get executed exactly 1️⃣ time.

func task() {
  fmt.Println("task done")
}

My first thought was channels. Create a channel, run the task in a goroutine and let the task wait till someone adds a value to the channel. Pass the channel to all goroutines that can signal the task to run.

The goroutines now just write to the channel once to signal. The code had this structure.

func parent() {
  doTask := make(chan bool)
  go func() {
    <- doTask
    task()      // runs the task once
  }
  // pass the channel around to goroutines
  go func(task) {
    ...
    if <some condition > {
      doTask <- true 
    }
    ...
  }
}

But now two problems arise:

1. Who closes the channel (channels can be closed only once).
2. How do we know if the task has already been done.

Solution to 1: Who closes the channel

The obvious initial answer to Who closes the channel seems to be the goroutine that creates the channel. Something like this:

func parent() {
  doTask := make(chan bool)
  defer close(doTask)
  go func() {
    <- doTask
    task()      // runs the task once
  }
// pass the channel around to goroutines
  go func(task) {
    ...
    if <some condition > {
      doTask <- true 
    }
    ...
  }
}

There’s an issue here if the parent goroutine completes first. The solution is to use a waitgroup and pass it to all goroutines. Then wait inside the parent for all goroutines to complete.

func parent() {
  doTask := make(chan bool)
  defer close(doTask)
  wg := sync.WaitGroup{}
  go func() {
    <- doTask
    task()      // runs the task once
  }
// pass the channel around to goroutines
  go func(task, wg) {
    wg.Add()    
    defer wg.Done()
    ...
    if <some condition > {
      doTask <- true 
    }
    ...
  }
  wg.Wait()
}

This solves our first problem. Atleast looks like.

The doTask channel is an unbuffered channel. So any goroutine that writes to the channel gets blocked till someone doesn’t read from it. If there are multipel goroutines which can write to it, all the goroutines other then the first one would get blocked. There’s a simple solution to that.

func parent() {
  doTask := make(chan bool)
  defer close(doTask)
  wg := sync.WaitGroup{}
  go func() {
    <- doTask
    task()      // runs the task once
    for _ := range doTask {} // drain till doTask gets closed
  }
  // pass the channel around to goroutines
  go func(task, wg) {
    wg.Add()    
    defer wg.Done()
    ...
    if <some condition > {
      doTask <- true 
    }
    ...
  }
  wg.Wait()
}

This finally solves our first issue ✅.

Solution to 2: Identify if the task is done

This was not initially needed, and hence the previous code kept working happily till the day we needed to know if the task was done and then do something.

With the channel thing we were sending information/signal one way. The task gets to know when it needs to run.

One option was to set a flag kind of thing and set it in the task before the loop to drain the doTask channel. And then in the goroutines check if it has been set.

To do so, I would need to pass down my flag variable as a pointer and missing to do it would be catastrophic 🤯.

So I went searching online for something like “golang run code once”. And I found this. The sync package contains something called as Once. I had found the perfect solution. I went experimenting and a working snippet looked like this

package main
import (
  "fmt"
  "sync"
  "time"
)
func main() {
  var once sync.Once
  
  task := func() {
    fmt.Println("Only once")
  }
 
  goroutine(&once, task)
  goroutine(&once, task)
 
  time.Sleep(1) // wait for goroutines to complete
}
func goroutine(once *sync.Once, onceBody func()) {
  once.Do(onceBody)
}

This worked but again how do I know if the job is done. Not possible !!

Another issue is if you pass the once by value instead of reference, things get messed up. But thankfully, vet reports these cases nicely with a message like call of goroutine copies lock value: sync.Once contains sync.Mutex

So I thought of implementing my own stateful once with a Done() method. First thing was to go and look at the source of Once. And after removing comments, it’s just this:

package sync
import (
  "sync/atomic"
)
// A Once must not be copied after first use.
type Once struct {
  done uint32
  m    Mutex
}
func (o *Once) Do(f func()) {
  if atomic.LoadUint32(&o.done) == 0 {
    o.doSlow(f)
  }
}
func (o *Once) doSlow(f func()) {
  o.m.Lock()
  defer o.m.Unlock()
  if o.done == 0 {
    defer atomic.StoreUint32(&o.done, 1)
    f()
  }
}

Thanks to the go team for keeping it open and easily accessible.
Now it’s time to add a Done() option. Doing that is simple. Just add another Done() that checks if done is set to 1 or not.

package sync
import (
  "sync/atomic"
)
// A Once must not be copied after first use.
type Once struct {
  done uint32
  m    Mutex
}
func (o *Once) Do(f func()) {
  if atomic.LoadUint32(&o.done) == 0 {
    o.doSlow(f)
  }
}
func (o *Once) doSlow(f func()) {
  o.m.Lock()
  defer o.m.Unlock()
  if o.done == 0 {
    defer atomic.StoreUint32(&o.done, 1)
    f()
  }
}
func (o *Once) Done() bool {
  return atomic.LoadUint32(&o.done) == 1
}

But then I realised there’s another problem here 🧐. Each goroutine now needs two things to be passed, the function body and the once object. As in the experiment example.

Final Solution

The last remaining thing was to have a function pointer within the Once struct, so that the function gets carried along with Once objects.

package sync
import (
  "sync/atomic"
)
// A Once must not be copied after first use.
type Once struct {
  f    func()
  done uint32
  m    Mutex
}
func (o *Once) Do() {
  if atomic.LoadUint32(&o.done) == 0 {
    o.doSlow()
  }
}
func (o *Once) doSlow() {
  o.m.Lock()
  defer o.m.Unlock()
  if o.done == 0 {
    defer atomic.StoreUint32(&o.done, 1)
    o.f()
  }
}
func (o *Once) Done() bool {
  return atomic.LoadUint32(&o.done) == 1
}

And now my example looks like this:

func main() {
  var once = Once {
    func() {
      fmt.Println("Only once")
    },
    0,
    sync.Mutex{},
  }
  
  goroutine(&once)
  goroutine(&once)
 
  time.Sleep(1) // wait for goroutines to complete
}
func goroutine(once *Once) {
  once.Do()
}

That where the story ends. Feel free to use above code. If you need something more, the package is available on pkg.go.dev. This package has couple more options and is tested for data races. Source available here. Or just go get github.com/leangaurav/sync.

You can find me on Linkedin 👋🏻👋🏻.