ctxio: A new Go package for cancellable I/O

I wanted to share a new open source Go package I wrote: jayconrod.com/ctxio. It's a small library that lets you cancel long-running I/O operations using context.Context.

Let's say you're writing a command-line tool that copies large files, perhaps on a slow network file system. When the user presses ^C, you want to cancel ongoing copy operations and clean up partially written files. Or perhaps, while handling an HTTP request, you need to communicate with another server using a raw net.Conn. If your HTTP request is cancelled (the client disconnects), you want to close the net.Conn, too, without waiting for a timeout.

The normal Go way to cancel a long-running operation is with a context. For example, if you want ^C to cancel something, you can use signal.NotifyContext, which returns a context that will be cancelled when the program receives any of the listed signals (os.Interrupt corresponds to SIGINT, the signal the terminal sends when ^C is pressed). Server libraries like net/http also provide a context that is cancelled if the client disconnects or times out. You can check whether a context is cancelled by calling context.Context.Err() or receiving from the channel returned by context.Context.Done().

Unfortunately, this doesn't help you with long-running I/O calls like io.Copy, which calls Read on a source and Write on a destination in a loop until there's nothing left to copy. Read and Write can block indefinitely, and even if they're fast, there's nowhere to plug in the context to stop the operation.

As an example, below we have a function that copies a source file to a destination, given two file paths. It uses io.Copy, and there's no way to interrupt it.

func CopyFile(dst, src string) (err error) {
  w, err := os.Create(dst)
  if err != nil {
    return err
  }
  defer func() {
    if cerr := w.Close(); err != nil {
      err = cerr
    }
  }()
  r, err := os.Open(src)
  if err != nil {
    return err
  }
  defer r.Close()
  _, err = io.Copy(w, r)   // <-- Might run indefinitely.
  return err
}

For many implementations, you can interrupt a Read or a Write by calling the Close method to close the file or shutdown the network connection. Close causes pending Read and Write operations to return immediately (perhaps with partial data), and future Read and Write operations will fail. This isn't guaranteed for all implementations by the contract for io.Reader, io.Writer, or io.Closer, but it works for *os.File and net.Conn, which is mostly what I care about.

Building on our example, let's add a "watchdog" goroutine that receives from ctx.Done(), then closes both files early.

func CopyFile(ctx context.Context, dst, src string) (err error) {
  // ... open files as before ...
  doneC := make(chan struct{})
  defer close(doneC)
  go func() {
    select {
    case <-ctx.Done():
      w.Close()
      r.Close()
    }
    case <-doneC:
  }()
  _, err = io.Copy(w, r)
  if ctx.Err() != nil {
    return ctx.Err()
  }
  return err
}

This adds a lot of boilerplate around a single io.Copy call and is tricky to get right. Which is why I wrote ctxio to do it for you. ctxio has three functions:

func Copy(ctx context.Context, dst io.WriteCloser, src io.ReadCloser) (written int64, err error)
func CopyN(ctx context.Context, dst io.WriteCloser, src io.ReadCloser, n int64) (written int64, err error)
func ReadAll(ctx context.Context, src io.ReadCloser) ([]byte, error)

These functions are replacements for io.Copy, io.CopyN, and io.ReadAll, respectively. Each accepts a context.Context as its first parameter. Instead of accepting io.Reader and io.Writer, these functions accept io.ReadCloser and io.WriteCloser. If the context is cancelled during a copy operation, Copy and friends call the Close method on both src and dst, which should cause any pending Read and Write calls to finish quickly.

Using ctxio.Copy, we can shorten our example:

func CopyFile(ctx context.Context, dst, src string) (err error) {
  // ... open files as before ...
  return ctxio.Copy(ctx, w, r)
}

If you're mainly worried about cancelling long-running transfers via net/http, you probably don't need this package unless you need to add your own cancellation mechanisms beyond what net/http provides. The package documentation doesn't say much about cancellation, but if you create requests with http.NewRequestWithContext on the client side, and you respect http.Request.Context on the server side, you're probably fine. The convenience functions like http.Get don't accept a context, so avoid them if you want to support cancellation.