github.com/pdmccormick/importable-docker-buildx@v0.0.0-20240426161518-e47091289030/util/ioset/ioset.go (about) 1 package ioset 2 3 import ( 4 "io" 5 "sync" 6 7 "github.com/pkg/errors" 8 "github.com/sirupsen/logrus" 9 ) 10 11 // Pipe returns a pair of piped readers and writers collection. 12 // They are useful for controlling stdio stream using Forwarder function. 13 func Pipe() (In, Out) { 14 r1, w1 := io.Pipe() 15 r2, w2 := io.Pipe() 16 r3, w3 := io.Pipe() 17 return In{r1, w2, w3}, Out{w1, r2, r3} 18 } 19 20 type In struct { 21 Stdin io.ReadCloser 22 Stdout io.WriteCloser 23 Stderr io.WriteCloser 24 } 25 26 func (s In) Close() (retErr error) { 27 if err := s.Stdin.Close(); err != nil { 28 retErr = err 29 } 30 if err := s.Stdout.Close(); err != nil { 31 retErr = err 32 } 33 if err := s.Stderr.Close(); err != nil { 34 retErr = err 35 } 36 return 37 } 38 39 type Out struct { 40 Stdin io.WriteCloser 41 Stdout io.ReadCloser 42 Stderr io.ReadCloser 43 } 44 45 func (s Out) Close() (retErr error) { 46 if err := s.Stdin.Close(); err != nil { 47 retErr = err 48 } 49 if err := s.Stdout.Close(); err != nil { 50 retErr = err 51 } 52 if err := s.Stderr.Close(); err != nil { 53 retErr = err 54 } 55 return 56 } 57 58 // Forwarder forwards IO between readers and writers contained 59 // in In and Out structs. 60 // In and Out can be changed during forwarding using SetIn and SetOut methods. 61 type Forwarder struct { 62 stdin *SingleForwarder 63 stdout *SingleForwarder 64 stderr *SingleForwarder 65 mu sync.Mutex 66 67 // PropagateStdinClose indicates whether EOF from Stdin of Out should be propagated. 68 // If this is true, EOF from Stdin (reader) of Out closes Stdin (writer) of In. 69 PropagateStdinClose bool 70 } 71 72 func NewForwarder() *Forwarder { 73 return &Forwarder{ 74 stdin: NewSingleForwarder(), 75 stdout: NewSingleForwarder(), 76 stderr: NewSingleForwarder(), 77 PropagateStdinClose: true, 78 } 79 } 80 81 func (f *Forwarder) Close() (retErr error) { 82 if err := f.stdin.Close(); err != nil { 83 retErr = err 84 } 85 if err := f.stdout.Close(); err != nil { 86 retErr = err 87 } 88 if err := f.stderr.Close(); err != nil { 89 retErr = err 90 } 91 return retErr 92 } 93 94 func (f *Forwarder) SetOut(out *Out) { 95 f.mu.Lock() 96 if out == nil { 97 f.stdin.SetWriter(nil, func() io.WriteCloser { return nil }) 98 f.stdout.SetReader(nil) 99 f.stderr.SetReader(nil) 100 } else { 101 f.stdin.SetWriter(out.Stdin, func() io.WriteCloser { 102 if f.PropagateStdinClose { 103 out.Stdin.Close() // propagate EOF 104 logrus.Debug("forwarder: propagating stdin close") 105 return nil 106 } 107 return out.Stdin 108 }) 109 f.stdout.SetReader(out.Stdout) 110 f.stderr.SetReader(out.Stderr) 111 } 112 f.mu.Unlock() 113 } 114 115 func (f *Forwarder) SetIn(in *In) { 116 f.mu.Lock() 117 if in == nil { 118 f.stdin.SetReader(nil) 119 f.stdout.SetWriter(nil, func() io.WriteCloser { return nil }) 120 f.stderr.SetWriter(nil, func() io.WriteCloser { return nil }) 121 } else { 122 f.stdin.SetReader(in.Stdin) 123 f.stdout.SetWriter(in.Stdout, func() io.WriteCloser { 124 return in.Stdout // continue write; TODO: make it configurable if needed 125 }) 126 f.stderr.SetWriter(in.Stderr, func() io.WriteCloser { 127 return in.Stderr // continue write; TODO: make it configurable if needed 128 }) 129 } 130 f.mu.Unlock() 131 } 132 133 // SingleForwarder forwards IO from a reader to a writer. 134 // The reader and writer can be changed during forwarding 135 // using SetReader and SetWriter methods. 136 type SingleForwarder struct { 137 curR io.ReadCloser // closed when set another reader 138 curRMu sync.Mutex 139 curW io.WriteCloser // closed when set another writer 140 curWEOFHandler func() io.WriteCloser 141 curWMu sync.Mutex 142 143 updateRCh chan io.ReadCloser 144 doneCh chan struct{} 145 146 closeOnce sync.Once 147 } 148 149 func NewSingleForwarder() *SingleForwarder { 150 f := &SingleForwarder{ 151 updateRCh: make(chan io.ReadCloser), 152 doneCh: make(chan struct{}), 153 } 154 go f.doForward() 155 return f 156 } 157 158 func (f *SingleForwarder) doForward() { 159 var r io.ReadCloser 160 for { 161 readerInvalid := false 162 var readerInvalidMu sync.Mutex 163 copyReaderToWriter := false 164 if r != nil { 165 copyReaderToWriter = true 166 } 167 if copyReaderToWriter { 168 srcR := r 169 go func() { 170 buf := make([]byte, 4096) 171 readerClosed := false 172 for { 173 n, readErr := srcR.Read(buf) 174 if readErr != nil { 175 srcR.Close() 176 readerClosed = true 177 if !errors.Is(readErr, io.EOF) && !errors.Is(readErr, io.ErrClosedPipe) { 178 logrus.Debugf("single forwarder: reader error: %v", readErr) 179 return 180 } 181 } 182 183 f.curWMu.Lock() 184 w := f.curW 185 f.curWMu.Unlock() 186 if w != nil { 187 if _, err := w.Write(buf[:n]); err != nil && !errors.Is(err, io.ErrClosedPipe) { 188 logrus.Debugf("single forwarder: writer error: %v", err) 189 } 190 } 191 readerInvalidMu.Lock() 192 ri := readerInvalid 193 readerInvalidMu.Unlock() 194 if ri || readerClosed { 195 return 196 } 197 if readErr != io.EOF { 198 logrus.Debugf("unknown error: %v\n", readErr) 199 continue 200 } 201 202 f.curWMu.Lock() 203 var newW io.WriteCloser 204 if f.curWEOFHandler != nil { 205 newW = f.curWEOFHandler() 206 } 207 f.curW = newW 208 f.curWMu.Unlock() 209 return 210 } 211 }() 212 } 213 select { 214 case newR := <-f.updateRCh: 215 f.curRMu.Lock() 216 if f.curR != nil { 217 f.curR.Close() 218 } 219 f.curR = newR 220 r = newR 221 readerInvalidMu.Lock() 222 readerInvalid = true 223 readerInvalidMu.Unlock() 224 f.curRMu.Unlock() 225 case <-f.doneCh: 226 return 227 } 228 } 229 } 230 231 // Close closes the both of registered reader and writer and finishes the forwarder. 232 func (f *SingleForwarder) Close() (retErr error) { 233 f.closeOnce.Do(func() { 234 f.curRMu.Lock() 235 r := f.curR 236 f.curR = nil 237 f.curRMu.Unlock() 238 if r != nil { 239 if err := r.Close(); err != nil { 240 retErr = err 241 } 242 } 243 // TODO: Wait until read data fully written to the current writer if needed. 244 f.curWMu.Lock() 245 w := f.curW 246 f.curW = nil 247 f.curWMu.Unlock() 248 if w != nil { 249 if err := w.Close(); err != nil { 250 retErr = err 251 } 252 } 253 close(f.doneCh) 254 }) 255 return retErr 256 } 257 258 // SetWriter sets the specified writer as the forward destination. 259 // If curWEOFHandler isn't nil, this will be called when the current reader returns EOF. 260 func (f *SingleForwarder) SetWriter(w io.WriteCloser, curWEOFHandler func() io.WriteCloser) { 261 f.curWMu.Lock() 262 if f.curW != nil { 263 // close all stream on the current IO no to mix with the new IO 264 f.curW.Close() 265 } 266 f.curW = w 267 f.curWEOFHandler = curWEOFHandler 268 f.curWMu.Unlock() 269 } 270 271 // SetWriter sets the specified reader as the forward source. 272 func (f *SingleForwarder) SetReader(r io.ReadCloser) { 273 f.updateRCh <- r 274 }