Force priority of go select statement

Question

I have the following piece of code:

func sendRegularHeartbeats(ctx context.Context) {
    for {
        select {
        case <-ctx.Done():
            return
        case <-time.After(1 * time.Second):
            sendHeartbeat()
        }
    }
}

This function is executed in a dedicated go-routine and sends a heartbeat-message every second. The whole process should stop immediately when the context is canceled.

Now consider the following scenario:

ctx, cancel := context.WithCancel(context.Background())
cancel()
go sendRegularHeartbeats(ctx)

This starts the heartbeat-routine with a closed context. In such a case, I don't want any heartbeats to be transmitted. So the first case block in the select should be entered immediately.

However, it seems that the order in which case blocks are evaluated is not guaranteed, and that the code sometimes sends a heartbeat message, even though the context is already canceled.

What is the correct way to implement such a behaviour?

I could add a "isContextclosed"-check in the second case, but that looks more like an ugly workaround for the problem.

Answer 1

The accepted answer has a wrong suggestion:

func sendRegularHeartbeats(ctx context.Context) {
    ticker := time.NewTicker(time.Second)
    defer ticker.Stop()

    for {
        //first select 
        select {
        case <-ctx.Done():
            return
        default:
        }

        //second select
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            sendHeartbeat()
        }
    }
}

This doesn't help, because of the following scenario:

1. both channels are empty
2. first select runs
3. both channels get a message concurrently
4. you are in the same probability game as if you haven't done anything in the first select

An alternative but still imperfect way is to guard against concurrent Done() events (the "wrong select") after consuming the ticker event i.e.

func sendRegularHeartbeats(ctx context.Context) {
    ticker := time.NewTicker(time.Second)
    defer ticker.Stop()

    for {            
        //select as usual
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            //give priority to a possible concurrent Done() event non-blocking way
            select {
              case <-ctx.Done():
              return
            default:
            }
            sendHeartbeat()
        }
    }
}

Caveat: the problem with this one is that it allows for "close enough" events to be confused - e.g. even though a ticker event arrived earlier, the Done event came soon enough to preempt the heartbeat. There is no perfect solution as of now.

Answer 2

Note beforehand:

Your example will work as you intend it to, as if the context is already cancelled when sendRegularHeartbeats() is called, the case <-ctx.Done() communication will be the only one ready to proceed and therefore chosen. The other case <-time.After(1 * time.Second) will only be ready to proceed after 1 second, so it will not be chosen at first. But to explicitly handle priorities when multiple cases might be ready, read on.

---

Unlike the case branches of a switch statement (where the evaluation order is the order they are listed), there is no priority or any order guaranteed in the case branches of a select statement.

Quoting from Spec: Select statements:

If one or more of the communications can proceed, a single one that can proceed is chosen via a uniform pseudo-random selection. Otherwise, if there is a default case, that case is chosen. If there is no default case, the "select" statement blocks until at least one of the communications can proceed.

If more communications can proceed, one is chosen randomly. Period.

If you want to maintain priority, you have to do that yourself (manually). You may do it using multiple select statements (subsequent, not nested), listing ones with higher priority in an earlier select, also be sure to add a default branch to avoid blocking if those are not ready to proceed. Your example requires 2 select statements, first one checking <-ctx.Done() as that is the one you want higher priority for.

I also recommend using a single time.Ticker instead of calling time.After() in each iteration (time.After() also uses a time.Ticker under the hood, but it doesn't reuse it just "throws it away" and creates a new one on the next call).

Here's an example implementation:

func sendRegularHeartbeats(ctx context.Context) {
    ticker := time.NewTicker(time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-ctx.Done():
            return
        default:
        }

        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            sendHeartbeat()
        }
    }
}

This will send no heartbeat if the context is already cancelled when sendRegularHeartbeats() is called, as you can check / verify it on the Go Playground.

If you delay the cancel() call for 2.5 seconds, then exactly 2 heartbeats will be sent:

ctx, cancel := context.WithCancel(context.Background())
go sendRegularHeartbeats(ctx)
time.Sleep(time.Millisecond * 2500)
cancel()
time.Sleep(time.Second * 2)

Try this one on the Go Playground.