]>
Commit | Line | Data |
---|---|---|
9b12e4fe JC |
1 | package ioutils |
2 | ||
3 | import ( | |
4 | "errors" | |
5 | "io" | |
6 | "net/http" | |
7 | "sync" | |
8 | ) | |
9 | ||
10 | // WriteFlusher wraps the Write and Flush operation ensuring that every write | |
11 | // is a flush. In addition, the Close method can be called to intercept | |
12 | // Read/Write calls if the targets lifecycle has already ended. | |
13 | type WriteFlusher struct { | |
14 | mu sync.Mutex | |
15 | w io.Writer | |
16 | flusher http.Flusher | |
17 | flushed bool | |
18 | closed error | |
19 | ||
20 | // TODO(stevvooe): Use channel for closed instead, remove mutex. Using a | |
21 | // channel will allow one to properly order the operations. | |
22 | } | |
23 | ||
24 | var errWriteFlusherClosed = errors.New("writeflusher: closed") | |
25 | ||
26 | func (wf *WriteFlusher) Write(b []byte) (n int, err error) { | |
27 | wf.mu.Lock() | |
28 | defer wf.mu.Unlock() | |
29 | if wf.closed != nil { | |
30 | return 0, wf.closed | |
31 | } | |
32 | ||
33 | n, err = wf.w.Write(b) | |
34 | wf.flush() // every write is a flush. | |
35 | return n, err | |
36 | } | |
37 | ||
38 | // Flush the stream immediately. | |
39 | func (wf *WriteFlusher) Flush() { | |
40 | wf.mu.Lock() | |
41 | defer wf.mu.Unlock() | |
42 | ||
43 | wf.flush() | |
44 | } | |
45 | ||
46 | // flush the stream immediately without taking a lock. Used internally. | |
47 | func (wf *WriteFlusher) flush() { | |
48 | if wf.closed != nil { | |
49 | return | |
50 | } | |
51 | ||
52 | wf.flushed = true | |
53 | wf.flusher.Flush() | |
54 | } | |
55 | ||
56 | // Flushed returns the state of flushed. | |
57 | // If it's flushed, return true, or else it return false. | |
58 | func (wf *WriteFlusher) Flushed() bool { | |
59 | // BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to | |
60 | // be used to detect whether or a response code has been issued or not. | |
61 | // Another hook should be used instead. | |
62 | wf.mu.Lock() | |
63 | defer wf.mu.Unlock() | |
64 | ||
65 | return wf.flushed | |
66 | } | |
67 | ||
68 | // Close closes the write flusher, disallowing any further writes to the | |
69 | // target. After the flusher is closed, all calls to write or flush will | |
70 | // result in an error. | |
71 | func (wf *WriteFlusher) Close() error { | |
72 | wf.mu.Lock() | |
73 | defer wf.mu.Unlock() | |
74 | ||
75 | if wf.closed != nil { | |
76 | return wf.closed | |
77 | } | |
78 | ||
79 | wf.closed = errWriteFlusherClosed | |
80 | return nil | |
81 | } | |
82 | ||
83 | // NewWriteFlusher returns a new WriteFlusher. | |
84 | func NewWriteFlusher(w io.Writer) *WriteFlusher { | |
85 | var flusher http.Flusher | |
86 | if f, ok := w.(http.Flusher); ok { | |
87 | flusher = f | |
88 | } else { | |
89 | flusher = &NopFlusher{} | |
90 | } | |
91 | return &WriteFlusher{w: w, flusher: flusher} | |
92 | } |