Supposed that I'm writing an http handler, that do something else before returning a response, do I have to setup a listener to check wether the http request context has been canceled? so that it can return immediately, or is there any other way to exit the handler when the request context cancelled?

func handleSomething(w http.ResponseWriter, r *http.Request) {
    done := make(chan error)

    go func() {
        if err := doSomething(r.Context()); err != nil {
            done <- err
                        return
        }

        done <- nil
    }()

    select {
    case <-r.Context().Done():
        http.Error(w, r.Context().Err().Error(), http.StatusInternalServerError)
        return
    case err := <-done:
        if err != nil {
            http.Error(w, err.Error(), http.StatusInternalServerError)
            return
        }

        w.WriteHeader(http.StatusOK)
        w.Write([]byte("ok"))
    }
}

func doSomething(ctx context.Context) error {
    // simulate doing something for 1 second.
    time.Sleep(time.Second)
    return nil
}

I tried making a test for it, but after the context got cancelled, doSomething function didn't stop and still running in the background.

func TestHandler(t *testing.T) {
    mux := http.NewServeMux()
    mux.HandleFunc("/something", handleSomething)

    srv := http.Server{
        Addr:    ":8989",
        Handler: mux,
    }

    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        defer wg.Done()
        if err := srv.ListenAndServe(); err != nil {
            log.Println(err)
        }
    }()

    time.Sleep(time.Second)

    req, err := http.NewRequest(http.MethodGet, "http://localhost:8989/something", nil)
    if err != nil {
        t.Fatal(err)
    }

    cl := http.Client{
        Timeout: 3 * time.Second,
    }

    res, err := cl.Do(req)
    if err != nil {
        t.Logf("error: %s", err.Error())
    } else {
        t.Logf("request is done with status code %d", res.StatusCode)
    }

    go func() {
        <-time.After(10 * time.Second)
        shutdown, cancel := context.WithTimeout(context.Background(), 10*time.Second)
        defer cancel()

        srv.Shutdown(shutdown)
    }()

    wg.Wait()
}

func handleSomething(w http.ResponseWriter, r *http.Request) {
    done := make(chan error)

    go func() {
        if err := doSomething(r.Context()); err != nil {
            log.Println(err)
            done <- err
        }

        done <- nil
    }()

    select {
    case <-r.Context().Done():
        log.Println("context is done!")
        return
    case err := <-done:
        if err != nil {
            http.Error(w, err.Error(), http.StatusInternalServerError)
            return
        }

        w.WriteHeader(http.StatusOK)
        w.Write([]byte("ok"))
    }
}

func doSomething(ctx context.Context) error {
    return runInContext(ctx, func() {
        log.Println("doing something")
        defer log.Println("done doing something")

        time.Sleep(10 * time.Second)
    })
}

func runInContext(ctx context.Context, fn func()) error {
    ch := make(chan struct{})
    go func() {
        defer close(ch)
        fn()
    }()

    select {
    case <-ctx.Done():
        return ctx.Err()
    case <-ch:
        return nil
    }
}

CodePudding user response：

I've just refactored the solution provided a little bit and now it should work. Let me guide you through the relevant changes.

The doSomething function

func doSomething(ctx context.Context) error {
    fmt.Printf("%v - doSomething: start\n", time.Now())
    select {
    case <-ctx.Done():
        fmt.Printf("%v - doSomething: cancelled\n", time.Now())
        return ctx.Err()
    case <-time.After(3 * time.Second):
        fmt.Printf("%v - doSomething: processed\n", time.Now())
        return nil
    }
}

It waits for a cancellation input or after a delay of 3 seconds it returns to the caller. It accepts a context to listen for.

The handleSomething function

func handleSomething(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()

    fmt.Printf("%v - handleRequestCtx: start\n", time.Now())

    done := make(chan error)
    go func() {
        if err := doSomething(ctx); err != nil {
            fmt.Printf("%v - handleRequestCtx: error %v\n", time.Now(), err)
            done <- err
        }

        done <- nil
    }()

    select {
    case <-ctx.Done():
        fmt.Printf("%v - handleRequestCtx: cancelled\n", time.Now())
        return
    case err := <-done:
        if err != nil {
            fmt.Printf("%v - handleRequestCtx: error: %v\n", time.Now(), err)
            w.WriteHeader(http.StatusInternalServerError)
            return
        }
        fmt.Printf("%v - handleRequestCtx: processed\n", time.Now())
    }
}

Here, the logic is very similar to yours. In the select, we check whether the received error is nil or not, and based on this we return to the proper HTTP status code to the caller. If we receive a cancellation input, we cancel all the context chain.

The TestHandler function

func TestHandler(t *testing.T) {
    r := mux.NewRouter()
    r.HandleFunc("/demo", handleSomething)

    srv := http.Server{
        Addr:    ":8000",
        Handler: r,
    }

    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        defer wg.Done()
        if err := srv.ListenAndServe(); err != nil {
            fmt.Println(err.Error())
        }
    }()

    ctx := context.Background()
    ctx, cancel := context.WithTimeout(ctx, 1*time.Second) // request canceled
    // ctx, cancel := context.WithTimeout(ctx, 5*time.Second) // request processed
    defer cancel()

    req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost:8000/demo", nil)

    client := http.Client{}
    res, err := client.Do(req)
    if err != nil {
        fmt.Println(err.Error())
    } else {
        fmt.Printf("res status code: %d\n", res.StatusCode)
    }
    srv.Shutdown(ctx)

    wg.Wait()
}

Here, we spin up an HTTP server and issue an HTTP request to it through an http.Client. You can see that there are two statements to set the context timeout. If you use the one with the comment // request canceled, everything will be canceled, otherwise, if you use the other the request will be processed.
I Hope that this clarifies your question!