Search code examples
goasynchronoustimeoutgoroutinego-context

Is there a way to cancel a context after a delay after one goroutine returns?


Problem

Situation

I currently have a gin handler function that runs three separate queries in three separate goroutines using the same context. There is an err group ("golang.org/x/sync/errgroup") that uses this shared context and the handler waits for the err group before returning.

Objective

The behavior I am trying to implement is after one of the goroutines finishes, there should be a timeout enforced on the remaining goroutines, but also this context should also be cancelled if the gin request is cancelled (connection closed), meaning gin's ctx.Request.Context() would have to be used.

Potential Solutions

Current implementation

Currently, I have a context with timeout passed to an errgroup but this just enforces a timeout for all the goroutines.

timeoutCtx := context.WithTimeout(context.Background(), 10*time.Second)
g, err := errgroup.WithContext(timeoutCtx)

g.Go(func1)
g.Go(func2)
g.Go(func3)

err = g.Wait()

Using the gin request context is required so that if the connection is closed and the request is cancelled, the goroutines will also stop.

// ctx *gin.Context
g, err := errgroup.WithContext(ctx.Request.Context())

g.Go(func1)
g.Go(func2)
g.Go(func3)

err = g.Wait()

Using a channel to implement timeout

Source

package main

import (
    "fmt"
    "time"
)

func main() {

    c1 := make(chan string, 1)
    go func() {
        time.Sleep(2 * time.Second)
        c1 <- "result 1"
    }()

    select {
    case res := <-c1:
        fmt.Println(res)
    case <-time.After(1 * time.Second):
        fmt.Println("timeout 1")
    }

    c2 := make(chan string, 1)
    go func() {
        time.Sleep(2 * time.Second)
        c2 <- "result 2"
    }()
    select {
    case res := <-c2:
        fmt.Println(res)
    case <-time.After(3 * time.Second):
        fmt.Println("timeout 2")
    }
}

Combining channels and request context

This solution is close but not very elegant or complete.

cQueryDone := make(chan bool)

g, err := errgroup.WithContext(ctx.Request.Context())

g.Go(func1)
g.Go(func2)
g.Go(func3)

// assumes func1 func2 and func3 all have cQueryDone <- true

if <-cQueryDone {
    select {
        case <-cQueryDone:
            select {
                case <-cQueryDone:
                    // ctx.JSON
                    // return
                case <-time.After(1*time.Second):
                    // ctx.JSON
                    // return
            }
        case <-time.After(3*time.Second):
            // ctx.JSON
            // return
    }
}

err = g.Wait()

Is there a better and more idiomatic way to implement this behavior in Go?


Solution

  • Note that context.WithTimeout() :

    • can wrap any context (not just context.Background())
    • also returns a cancel function

    You can add a timeout on top of ctx.Request.Context(), and call cancel when any of the queries completes :

    timeoutCtx, cancel := context.WithTimeout(ctx.Request.Context())
    
    g, err := errgroup.WithContext(timeoutCtx)
    
    g.Go( func1(cancel) ) // pass the cancel callback to each query some way or another
    g.Go( func2(cancel) ) // you prabably want to also pass timeoutCtx
    g.Go( func3(cancel) )
    
    g.Wait()
    

    Following your comment : there is also context.WithCancel(), and you can call cancel after a delay

    childCtx, cancel := context.WithCancel(ctx.Request.Context())
    
    g, err := errgroup.WithContext(childCtx)
    
    hammerTime := func(){
        <-time.After(1*time.Second)
        cancel()
    }
    
    g.Go( func1(hammerTime) ) // funcXX should have access to hammerTime
    g.Go( func2(hammerTime) )
    g.Go( func3(hammerTime) )
    
    g.Wait()