src/io/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 }