github.com/haraldrudell/parl@v0.4.176/pio/context-writer.go (about)

     1  /*
     2  © 2023–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
     3  ISC License
     4  */
     5  
     6  package pio
     7  
     8  import (
     9  	"context"
    10  	"io"
    11  )
    12  
    13  // ContextWriter is an [io.WriteCloser] that aborts on context cancel
    14  //   - on context cancel, Write returns error [context.Canceled]
    15  //   - If the runtime type of writer implements [io.Close], it ContextWriter can close it
    16  type ContextWriter struct {
    17  	writer io.Writer // Write()
    18  	// idempotent pannic-free closer if reader implemented [io.Closer]
    19  	//	- Close() IsClosable()
    20  	ContextCloser
    21  	ctx context.Context
    22  }
    23  
    24  var _ io.WriteCloser = &ContextWriter{}
    25  
    26  // NewContextWriter returns an [io.WriteCloser] that aborts on context cancel
    27  //   - on context cancel, Write returns error [context.Canceled]
    28  //   - If the runtime type of reader implements [io.Close], it can be closed
    29  func NewContextWriter(writer io.Writer, ctx context.Context) (contextWriter *ContextWriter) {
    30  	var closer, _ = writer.(io.Closer)
    31  	return &ContextWriter{
    32  		writer:        writer,
    33  		ContextCloser: *NewContextCloser(closer),
    34  		ctx:           ctx,
    35  	}
    36  }
    37  
    38  // Write is like [io.Writer.Write] but cancels if the context is canceled
    39  //   - on context cancel, the error returned is [context.Canceled]
    40  func (c *ContextWriter) Write(p []byte) (n int, err error) {
    41  	if err = c.ctx.Err(); err != nil {
    42  		return
    43  	}
    44  	return c.writer.Write(p)
    45  }