CODEWITHSUNDEEP
gochannelcancellation
Rodrigo
Rodrigo

Reputation: 3278

Terminating function execution if a context is cancelled

I have this current function that originally was not context aware.

func (s *Service) ChunkUpload(r *multipart.Reader) error {
    chunk, err := s.parseChunk(r)
    if err != nil {
        return fmt.Errorf("failed parsing chunk %w", err)
    }

    if err := os.MkdirAll(chunk.UploadDir, 02750); err != nil {
        return err
    }

    if err := s.saveChunk(chunk); err != nil {
        return fmt.Errorf("failed saving chunk %w", err)
    }

    return nil
}

I've updated it's method call to now take a context.Context as its first argument. My main goal is to terminate and return the function as soon as the context is cancelled.

My initial implementation was this.

func (s *Service) ChunkUpload(ctx context.Context, r *multipart.Reader) error {
    errCh := make(chan error)

    go func() {
        chunk, err := s.parseChunk(r)
        if err != nil {
            errCh <- fmt.Errorf("failed parsing chunk %w", err)
            return
        }

        if err := os.MkdirAll(chunk.UploadDir, 02750); err != nil {
            errCh <- err
            return
        }

        if err := s.saveChunk(chunk); err != nil {
            errCh <- fmt.Errorf("failed saving chunk %w", err)
            return
        }
    }()

    select {
    case err := <-errCh:
        return err
    case <-ctx.Done():
        return ctx.Err()
    }
}

However, as I thought about the execution of the code I realized that this doesn't achieve my goal. Since all the function's logic is in a separate go routine even if context gets cancelled and I return ChunkUpload early the code within the go routine will continue to execute thus not really making a difference from the original code.

The next though was okay just pass a context to all inner functions like s.parseChunk and s.saveChunk but this option also doesn't seem right as I would need to implement cancellations in each function. What would be the proper way to refactor this original function to be context aware and terminate as soon as a context is cancelled?

Upvotes: 2

Views: 5414

Answers (2)

Billy Yuan
Billy Yuan

Reputation: 1305

When we talked about canceling, we talked about a long run function or a block repeating multiple times, such as http.Serve()

As to your case, assume saveChunk will cost seconds to run, and you want to cancel when it's saving. So we can split the chunk into pieces, and save one by one, after each piece.

for i:=0;i<n;i++{
    select {
        case err := <- s.saveChunk(chunk[i]):
        {
             if err != nil {
                  fmt.Errorf("failed saving chunk %w", err)
                  return
              }
        }
        case <-ctx.Done():
              return
    }
}

Upvotes: 0

icza
icza

Reputation: 418127

Function calls and goroutines cannot be terminated from the caller, the functions and goroutines have to support the cancellation, often via a context.Context value or a done channel.

In either case, the functions are responsible to check / monitor the context, and if cancel is requested (when the Context's done channel is closed), return early. There isn't an easier / automatic way.

If the task executes code in a loop, a convenient solution is to check the done channel in each iteration, and return if it's closed. If the task is one "monolith", the implementor is responsible to use / insert "checkpoints" at which the task can be reasonably aborted early if such cancellation is requested.

An easy way to check if the done channel is closed is to use a non-blocking select, such as:

select {
case <-ctx.Done():
    // Abort / return early
    return
default:
}

Care must be taken when the task uses other channel operations, as they may block in a nondeterministic way. Those selects should include the ctx.Done() channel too:

select {
case v := <- someChannel:
    // Do something with v
case <-ctx.Done():
    // Abort / return early
    return
}

Also be careful, because if the above receive from someChannel never blocks there is no guarantee a cancellation is properly handled, because if multiple communications can proceed in a select, one is chosen randomly (and there's no guarantee the <-ctx.Done() is ever chosen). In such case you may combine the above 2: first do a non-blocking check for cancellation, then use a select with your channel operations and the cancel monitoring.

Upvotes: 7

Related Questions