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  }