forked from dvyukov/go-fuzz
-
Notifications
You must be signed in to change notification settings - Fork 1
/
writerset.go
107 lines (90 loc) · 2.33 KB
/
writerset.go
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
// Copyright 2015 Dmitry Vyukov. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
// Package writerset implements a mechanism to add and remove writers from a construct
// similar to io.MultiWriter. Manually vendored from http://github.com/stephens2424/writerset
package writerset
import (
"io"
"net/http"
"sync"
)
// ErrPartialWrite encapsulates an error from a WriterSet.
type ErrPartialWrite struct {
Writer io.Writer
Err error
Expected, Wrote int
}
// Error returns the error string from the underlying error.
func (e ErrPartialWrite) Error() string {
return e.Err.Error()
}
// WriterSet wraps multiple writers like io.MultiWriter, but such that individual
// writers are easy to add or remove as necessary.
type WriterSet struct {
m map[io.Writer]chan error
mu sync.Mutex
}
// New initializes a new empty writer set.
func New() *WriterSet {
return &WriterSet{
m: make(map[io.Writer]chan error),
mu: sync.Mutex{},
}
}
// Add ensures w is in the set.
func (ws *WriterSet) Add(w io.Writer) <-chan error {
ws.mu.Lock()
defer ws.mu.Unlock()
c, ok := ws.m[w]
if ok {
return c
}
c = make(chan error, 1)
ws.m[w] = c
return c
}
// Contains determines if w is in the set.
func (ws *WriterSet) Contains(w io.Writer) bool {
ws.mu.Lock()
defer ws.mu.Unlock()
_, ok := ws.m[w]
return ok
}
// Remove ensures w is not in the set.
func (ws *WriterSet) Remove(w io.Writer) {
ws.mu.Lock()
defer ws.mu.Unlock()
delete(ws.m, w)
}
// Write writes data to each underlying writer. If an error occurs on an underlying writer,
// that writer is removed from the set. The error will be wrapped as an ErrPartialWrite and
// sent on the channel created when the writer was added.
func (ws *WriterSet) Write(b []byte) (int, error) {
ws.mu.Lock()
defer ws.mu.Unlock()
for w, c := range ws.m {
bs, err := w.Write(b)
if err != nil {
c <- ErrPartialWrite{
Err: err,
Wrote: bs,
Expected: len(b),
Writer: w,
}
close(c)
delete(ws.m, w)
}
}
return len(b), nil
}
// Flush implements http.Flusher by calling flush on all the underlying writers if they are
// also http.Flushers.
func (ws *WriterSet) Flush() {
ws.mu.Lock()
defer ws.mu.Unlock()
for w := range ws.m {
if w, ok := w.(http.Flusher); ok {
w.Flush()
}
}
}