Wait for multiple callbacks with timeout in go without busy waiting or polling

In go I have two callbacks that eventually do not fire.

registerCb(func() {...})
registerCb(func() {...})
/* Wait for both func to execute with timeout */

I want to wait for both of them but having a timeout if one is not executed.

sync.WaitGroup does not work, since it is blocking and not channel based. Also you call WaitGroup.Done() without the risk of panic outside the callbacks.

My current solution is using just two booleans and a busy wait loop. But that's not satisfying.
Is there any idiomatic way that do not use polling or busy waiting?

Update:

Here is some code that demonstrates a busy wait solution but should return as soon as both callbacks are fired or after the timeout, without using polling

package main
import (
"fmt"
"log"
"sync"
"time"
)
var cbOne func()
var cbTwo func()
func registerCbOne(cb func()) {
cbOne = cb
}
func registerCbTwo(cb func()) {
cbTwo = cb
}
func executeCallbacks() {
<-time.After(1 * time.Second)
cbOne()
// Might never happen
//<-time.After(1 * time.Second)
//cbTwo()
}
func main() {
// Some process in background will execute our callbacks
go func() {
executeCallbacks()
}()
err := WaitAllOrTimeout(3 * time.Second)
if err != nil {
fmt.Println("Error: ", err.Error())
}
fmt.Println("Hello, playground")
}
func WaitAllOrTimeout(to time.Duration) error {
cbOneDoneCh := make(chan bool, 1)
cbTwoDoneCh := make(chan bool, 1)
cbOneDone := false
cbTwoDone := false
registerCbOne(func() {
fmt.Println("cb One");
cbOneDoneCh <- true
})
registerCbTwo(func() {
fmt.Println("cb Two");
cbTwoDoneCh <- true
})
// Wait for cbOne and cbTwo to be executed or a timeout
// Busywait solution
for {
select {
case <-time.After(to):
if cbOneDone && cbTwoDone {
fmt.Println("Both CB executed (we could poll more often)")
return nil
}
fmt.Println("Timeout!")
return fmt.Errorf("Timeout")
case <-cbOneDoneCh:
cbOneDone = true
case <-cbTwoDoneCh:
cbTwoDone = true
}
}
}

func wait(ctx context.Context, wg *sync.WaitGroup) error {
done := make(chan struct{}, 1)
go func() {
wg.Wait()
done <- struct{}{}
}()
select {
case <-done:
// Counter is 0, so all callbacks completed.
return nil
case <-ctx.Done():
// Context cancelled.
return ctx.Err()
}
}

Alternatively, you can pass a time.Duration and block on <-time.After(d) rather than on <-ctx.Done(), but I would argue that using context is more idiomatic.

below code present two variations,

• the first is the regular pattern, nothing fancy, it does the job and does it well. You launch your callbacks into a routine, you make them push to a sink, listen that sink for a result or timeout. Take care to the sink channel initial capacity, to prevent leaking a routine it must match the number of callbacks.
• the second factories out the synchronization mechanisms into small functions to assemble, two wait methods are provided, waitAll and waitOne. Nice to write, but definitely less efficient, more allocations, more back and forth with more channels, more complex to reason about, more subtle.

package main

import (
	"fmt"
	"log"
	"sync"
	"time"
)

func main() {
	ExampleOne()
	ExampleTwo()
	ExampleThree()

	fmt.Println("Hello, playground")
}

func ExampleOne() {
	log.Println("start reg")
	errs := make(chan error, 2)
	go func() {
		fn := callbackWithOpts("reg: so slow", 2*time.Second, nil)
		errs <- fn()
	}()
	go func() {
		fn := callbackWithOpts("reg: too fast", time.Millisecond, fmt.Errorf("broke!"))
		errs <- fn()
	}()

	select {
	case err := <-errs: // capture only one result,
		// the fastest to finish.
		if err != nil {
			log.Println(err)
		}
	case <-time.After(time.Second): // or wait that many amount of time,
		// in case they are all so slow.
	}
	log.Println("done reg")
}

func ExampleTwo() {
	log.Println("start wait")
	errs := waitAll(
		withTimeout(time.Second,
			callbackWithOpts("waitAll: so slow", 2*time.Second, nil),
		),
		withTimeout(time.Second,
			callbackWithOpts("waitAll: too fast", time.Millisecond, nil),
		),
	)
	for err := range trim(errs) {
		if err != nil {
			log.Println(err)
		}
	}
	log.Println("done wait")
}

func ExampleThree() {
	log.Println("start waitOne")
	errs := waitOne(
		withTimeout(time.Second,
			callbackWithOpts("waitOne: so slow", 2*time.Second, nil),
		),
		withTimeout(time.Second,
			callbackWithOpts("waitOne: too fast", time.Millisecond, nil),
		),
	)
	for err := range trim(errs) {
		if err != nil {
			log.Println(err)
		}
	}
	log.Println("done waitOne")
}

// a configurable callback for playing
func callbackWithOpts(msg string, tout time.Duration, err error) func() error {
	return func() error {
		<-time.After(tout)
		fmt.Println(msg)
		return err
	}
}

// withTimeout return a function that returns first error or times out and return nil
func withTimeout(tout time.Duration, h func() error) func() error {
	return func() error {
		d := make(chan error, 1)
		go func() {
			d <- h()
		}()
		select {
		case err := <-d:
			return err
		case <-time.After(tout):
		}
		return nil
	}
}

// wait launches all func() and return their errors into the returned error channel; (merge)
// It is the caller responsability to drain the output error channel.
func waitAll(h ...func() error) chan error {
	d := make(chan error, len(h))
	var wg sync.WaitGroup
	for i := 0; i < len(h); i++ {
		wg.Add(1)
		go func(h func() error) {
			defer wg.Done()
			d <- h()
		}(h[i])
	}
	go func() {
		wg.Wait()
		close(d)
	}()
	return d
}

// wait launches all func() and return the first error into the returned error channel
// It is the caller responsability to drain the output error channel.
func waitOne(h ...func() error) chan error {
	d := make(chan error, len(h))
	one := make(chan error, 1)
	var wg sync.WaitGroup
	for i := 0; i < len(h); i++ {
		wg.Add(1)
		go func(h func() error) {
			defer wg.Done()
			d <- h()
		}(h[i])
	}
	go func() {
		for err := range d {
			one <- err
		        close(one)
			break
		}
	}()
	go func() {
		wg.Wait()
		close(d)
	}()
	return one
}

func trim(err chan error) chan error {
	out := make(chan error)
	go func() {
		for e := range err {
			out <- e
		}
		close(out)
	}()
	return out
}

This is a followup to my comment, added after you added your example solution. To be clearer than I can in comments, your example code is actually not that bad. Here is your original example:

> // Busywait solution
> for {
> select {
> case <-time.After(to):
> if cbOneDone && cbTwoDone {
> fmt.Println("Both CB executed (we could poll more often)")
> return nil
> }
> fmt.Println("Timeout!")
> return fmt.Errorf("Timeout")
> case <-cbOneDoneCh:
> cbOneDone = true
> case <-cbTwoDoneCh:
> cbTwoDone = true
> }
> }

This isn't a "busy wait" but it does have several bugs (including the fact that you need an only-once send semantic for the done channels, or maybe easier and at least as good, to just close them once when done, perhaps using sync.Once). What we want to do is:

1. Start a timer with to as the timeout.
2. Enter a select loop, using the timer's channel and the two "done" channels.

We want to exit the select loop when the first of the following events occurs:

• the timer fires, or
• both "done" channels have been signaled.

If we're going to close the two done channels we'll want to have the Ch variables cleared (set to nil) as well so that the selects don't spin—that would turn this into a true busy-wait—but for the moment let's just assume instead that we send exactly once on them on callback, and otherwise just leak the channels, so that we can use your code as written as those selects will only ever return once. Here's the updated code:

t := timer.NewTimer(to)
for !cbOneDone || !cbTwoDone {
select {
case <-t.C:
fmt.Println("Timeout!")
return fmt.Errorf("timeout")
}
case <-cbOneDoneCh:
cbOneDone = true
case <-cbTwoDoneCh:
cbTwoDone = true
}
}
// insert t.Stop() and receive here to drain t.C if desired
fmt.Println("Both CB executed")
return nil

Note that we will go through the loop at most two times:

• If we receive from both Done channels, once each, the loop stops without a timeout. There's no spinning/busy-waiting: we never received anything from t.C. We return nil (no error).

• If we receive from one Done channel, the loop resumes but blocks waiting for the timer or the other Done channel.

• If we ever receive from t.C, it means we didn't get both callbacks yet. We may have had one, but there's been a timeout and we choose to give up, which was our goal. We return an error, without going back through the loop.

A real version needs a bit more work to clean up properly and avoid leaking "done" channels (and the timer channel and its goroutine; see comment), but this is the general idea. You're already turning the callbacks into channel operations, and you already have a timer with its channel.