github.com/mutagen-io/mutagen@v0.18.0-rc1/pkg/stream/preemptable_writer.go (about)

     1  package stream
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  )
     7  
     8  var (
     9  	// ErrWritePreempted indicates that a write operation was preempted.
    10  	ErrWritePreempted = errors.New("write preempted")
    11  )
    12  
    13  // preemptableWriter is an io.Writer implementation that checks for preemption
    14  // every N writes.
    15  type preemptableWriter struct {
    16  	// writer is the underlying writer.
    17  	writer io.Writer
    18  	// cancelled is the channel that, when closed, indicates preemption.
    19  	cancelled <-chan struct{}
    20  	// checkInterval is the number of writes to allow between preemption checks.
    21  	checkInterval uint
    22  	// writeCount is the number of writes since the last preemption check.
    23  	writeCount uint
    24  }
    25  
    26  // NewPreemptableWriter wraps an io.Writer and provides preemption capabilities
    27  // for long copy operations. It takes an underlying writer, a channel that (once
    28  // closed) indicates cancellation, and an interval that specifies the maximum
    29  // number of Write calls that should be processed between cancellation checks.
    30  // If interval is 0, a cancellation check will be performed before every write.
    31  func NewPreemptableWriter(writer io.Writer, cancelled <-chan struct{}, interval uint) io.Writer {
    32  	return &preemptableWriter{
    33  		writer:        writer,
    34  		cancelled:     cancelled,
    35  		checkInterval: interval,
    36  	}
    37  }
    38  
    39  // Write implements io.Writer.Write.
    40  func (w *preemptableWriter) Write(data []byte) (int, error) {
    41  	// Handle preemption checking.
    42  	if w.writeCount == w.checkInterval {
    43  		select {
    44  		case <-w.cancelled:
    45  			return 0, ErrWritePreempted
    46  		default:
    47  		}
    48  		w.writeCount = 0
    49  	} else {
    50  		w.writeCount++
    51  	}
    52  
    53  	// Perform the write.
    54  	return w.writer.Write(data)
    55  }