blob: cb8df57c99ab0849d58d5a1ce34a66c3bbade0f4 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
package message
import (
"bytes"
"io"
"sync"
"sync/atomic"
"syscall"
)
const (
suspendBufInitial = 1 << 12
suspendBufMax = 1 << 24
)
// Suspendable proxies writes to a downstream [io.Writer] but optionally withholds writes
// between calls to Suspend and Resume.
type Suspendable struct {
Downstream io.Writer
s atomic.Bool
buf bytes.Buffer
// for growing buf
bufOnce sync.Once
// for synchronising all other buf operations
bufMu sync.Mutex
dropped int
}
func (s *Suspendable) Write(p []byte) (n int, err error) {
if !s.s.Load() {
return s.Downstream.Write(p)
}
s.bufOnce.Do(func() { s.buf.Grow(suspendBufInitial) })
s.bufMu.Lock()
defer s.bufMu.Unlock()
if free := suspendBufMax - s.buf.Len(); free < len(p) {
// fast path
if free <= 0 {
s.dropped += len(p)
return 0, syscall.ENOMEM
}
n, _ = s.buf.Write(p[:free])
err = syscall.ENOMEM
s.dropped += len(p) - n
return
}
return s.buf.Write(p)
}
// IsSuspended returns whether [Suspendable] is currently between a call to Suspend and Resume.
func (s *Suspendable) IsSuspended() bool { return s.s.Load() }
// Suspend causes [Suspendable] to start withholding output in its buffer.
func (s *Suspendable) Suspend() bool { return s.s.CompareAndSwap(false, true) }
// Resume undoes the effect of Suspend and dumps the buffered into the downstream [io.Writer].
func (s *Suspendable) Resume() (resumed bool, dropped uintptr, n int64, err error) {
if s.s.CompareAndSwap(true, false) {
s.bufMu.Lock()
defer s.bufMu.Unlock()
resumed = true
dropped = uintptr(s.dropped)
s.dropped = 0
n, err = io.Copy(s.Downstream, &s.buf)
s.buf.Reset()
}
return
}
|
