Writing while respecting Context.Done

writectx.go

// writeCtx writes bytes to a writer while making a best effort to
// respect the Done signal of the Context.  However, once any bytes have
// been written to w, writeCtx will ignore the Done signal to avoid
// partial writes.
func writeCtx(ctx context.Context, w io.Writer, b []byte) (int, error) {
	select {
	case <-ctx.Done():
		// Early cancel.
		return 0, ctx.Err()
	default:
	}
	// Check for timeout support.
	wd, ok := w.(interface {
		SetWriteDeadline(time.Time) error
	})
	if !ok {
		return w.Write(b)
	}
	if err := wd.SetWriteDeadline(time.Now()); err != nil {
		return w.Write(b)
	}
	// Start separate goroutine to wait on Context.Done.
	if d, ok := ctx.Deadline(); ok {
		wd.SetWriteDeadline(d)
	} else {
		wd.SetWriteDeadline(time.Time{})
	}
	writeDone := make(chan struct{})
	listenDone := make(chan struct{})
	go func() {
		defer close(listenDone)
		select {
		case <-ctx.Done():
			wd.SetWriteDeadline(time.Now()) // interrupt write
		case <-writeDone:
		}
	}()
	n, err := w.Write(b)
	close(writeDone)
	<-listenDone
	if n == 0 || !isTimeout(err) {
		return n, err
	}
	// Data has been written.  Block until finished, since partial writes
	// are guaranteed protocol violations.
	wd.SetWriteDeadline(time.Time{})
	nn, err := w.Write(b[n:])
	return n + nn, err
}

func isTimeout(e error) bool {
	te, ok := e.(interface {
		Timeout() bool
	})
	return ok && te.Timeout()
}