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 }