Wait for context Done channel for cancellation while working on long-run operation in Go

Issue

In this scenario:

  1. the 1st go routine is working on a long-run operation and block until it’s done
  2. the 2nd go routine may cancel the whole task at any time
  3. when the whole task is cancelled, the 1st go routine should just quit the operation and return immediately.

Here’s my solution. It’s works, but doesn’t feel elegant or in Go style.

Can you please correct me or show me a better solution?

    var (
        workTimeCost  = 6 * time.Second
        cancelTimeout = 5 * time.Second
    )

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

    var (
        data   int
        readCh = make(chan struct{})
    )
    go func() {
        log.Println("blocked to read data")
        // fake long i/o operations
        time.Sleep(workTimeCost)
        data = 10
        log.Println("done read data")

        readCh <- struct{}{}
    }()

    // fake cancel is called from the other routine (it's actually not caused by timeout)
    time.AfterFunc(cancelTimeout, cancel)

    select {
    case <-ctx.Done():
        log.Println("cancelled")
        return
    case <-readCh:
        break
    }

    log.Println("got final data", data)

Solution

Close readCh to indicate that the long running goroutine completed. There are two benefits to closing the channel compared to sending a value:

  • close is conveniently called with defer
  • close does not block in the case where the context was canceled. The code in the question leaks a goroutine if the context is canceled before the goroutine completes.

Here’s the updated code:

var (
    workTimeCost  = 6 * time.Second
    cancelTimeout = 5 * time.Second
)

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

var (
    data   int
    readCh = make(chan struct{})
)
go func() {
    defer close(readCh)
    log.Println("blocked to read data")
    // fake long i/o operations
    time.Sleep(workTimeCost)
    data = 10
    log.Println("done read data")
}()

// fake cancel is called from the other routine (it's actually not caused by timeout)
time.AfterFunc(cancelTimeout, cancel)

select {
case <-ctx.Done():
    log.Println("cancelled")
    return
case <-readCh:
    break
}

log.Println("got final data", data)

If you do not need to distinguish between completion of the long-running goroutine and cancelation, call the cancel function from the goroutine.

var (
    workTimeCost  = 6 * time.Second
    cancelTimeout = 5 * time.Second
)

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

var data int

go func() {
    defer cancel()
    log.Println("blocked to read data")
    // fake long i/o operations
    time.Sleep(workTimeCost)
    data = 10
    log.Println("done read data")
}()

// fake cancel is called from the other routine (it's actually not caused by timeout)
time.AfterFunc(cancelTimeout, cancel)

<-ctx.Done()

log.Println("got final data", data)

Answered By – Bayta Darell

Answer Checked By – Katrina (GoLangFix Volunteer)

Leave a Reply

Your email address will not be published.