multi.go

     1  // Copyright 2010 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package io
     6  
     7  type eofReader struct{}
     8  
     9  func (eofReader) Read([]byte) (int, error) {
    10  	return 0, EOF
    11  }
    12  
    13  type multiReader struct {
    14  	readers []Reader
    15  }
    16  
    17  func (mr *multiReader) Read(p []byte) (n int, err error) {
    18  	for len(mr.readers) > 0 {
    19  		// Optimization to flatten nested multiReaders (Issue 13558).
    20  		if len(mr.readers) == 1 {
    21  			if r, ok := mr.readers[0].(*multiReader); ok {
    22  				mr.readers = r.readers
    23  				continue
    24  			}
    25  		}
    26  		n, err = mr.readers[0].Read(p)
    27  		if err == EOF {
    28  			// Use eofReader instead of nil to avoid nil panic
    29  			// after performing flatten (Issue 18232).
    30  			mr.readers[0] = eofReader{} // permit earlier GC
    31  			mr.readers = mr.readers[1:]
    32  		}
    33  		if n > 0 || err != EOF {
    34  			if err == EOF && len(mr.readers) > 0 {
    35  				// Don't return EOF yet. More readers remain.
    36  				err = nil
    37  			}
    38  			return
    39  		}
    40  	}
    41  	return 0, EOF
    42  }
    43  
    44  func (mr *multiReader) WriteTo(w Writer) (sum int64, err error) {
    45  	return mr.writeToWithBuffer(w, make([]byte, 1024*32))
    46  }
    47  
    48  func (mr *multiReader) writeToWithBuffer(w Writer, buf []byte) (sum int64, err error) {
    49  	for i, r := range mr.readers {
    50  		var n int64
    51  		if subMr, ok := r.(*multiReader); ok { // reuse buffer with nested multiReaders
    52  			n, err = subMr.writeToWithBuffer(w, buf)
    53  		} else {
    54  			n, err = copyBuffer(w, r, buf)
    55  		}
    56  		sum += n
    57  		if err != nil {
    58  			mr.readers = mr.readers[i:] // permit resume / retry after error
    59  			return sum, err
    60  		}
    61  		mr.readers[i] = nil // permit early GC
    62  	}
    63  	mr.readers = nil
    64  	return sum, nil
    65  }
    66  
    67  var _ WriterTo = (*multiReader)(nil)
    68  
    69  // MultiReader returns a Reader that's the logical concatenation of
    70  // the provided input readers. They're read sequentially. Once all
    71  // inputs have returned EOF, Read will return EOF.  If any of the readers
    72  // return a non-nil, non-EOF error, Read will return that error.
    73  func MultiReader(readers ...Reader) Reader {
    74  	r := make([]Reader, len(readers))
    75  	copy(r, readers)
    76  	return &multiReader{r}
    77  }
    78  
    79  type multiWriter struct {
    80  	writers []Writer
    81  }
    82  
    83  func (t *multiWriter) Write(p []byte) (n int, err error) {
    84  	for _, w := range t.writers {
    85  		n, err = w.Write(p)
    86  		if err != nil {
    87  			return
    88  		}
    89  		if n != len(p) {
    90  			err = ErrShortWrite
    91  			return
    92  		}
    93  	}
    94  	return len(p), nil
    95  }
    96  
    97  var _ StringWriter = (*multiWriter)(nil)
    98  
    99  func (t *multiWriter) WriteString(s string) (n int, err error) {
   100  	var p []byte // lazily initialized if/when needed
   101  	for _, w := range t.writers {
   102  		if sw, ok := w.(StringWriter); ok {
   103  			n, err = sw.WriteString(s)
   104  		} else {
   105  			if p == nil {
   106  				p = []byte(s)
   107  			}
   108  			n, err = w.Write(p)
   109  		}
   110  		if err != nil {
   111  			return
   112  		}
   113  		if n != len(s) {
   114  			err = ErrShortWrite
   115  			return
   116  		}
   117  	}
   118  	return len(s), nil
   119  }
   120  
   121  // MultiWriter creates a writer that duplicates its writes to all the
   122  // provided writers, similar to the Unix tee(1) command.
   123  //
   124  // Each write is written to each listed writer, one at a time.
   125  // If a listed writer returns an error, that overall write operation
   126  // stops and returns the error; it does not continue down the list.
   127  func MultiWriter(writers ...Writer) Writer {
   128  	allWriters := make([]Writer, 0, len(writers))
   129  	for _, w := range writers {
   130  		if mw, ok := w.(*multiWriter); ok {
   131  			allWriters = append(allWriters, mw.writers...)
   132  		} else {
   133  			allWriters = append(allWriters, w)
   134  		}
   135  	}
   136  	return &multiWriter{allWriters}
   137  }
   138
View as plain text