github.com/jcarley/cli@v0.0.0-20180201210820-966d90434c30/lib/transfer/transfer.go (about)

     1  package transfer
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync/atomic"
     7  )
     8  
     9  type ByteSize float64
    10  
    11  const (
    12  	_           = iota
    13  	KB ByteSize = 1 << (10 * iota)
    14  	MB
    15  	GB
    16  	TB
    17  )
    18  
    19  func (b ByteSize) String() string {
    20  	switch {
    21  	case b >= TB:
    22  		return fmt.Sprintf("%.3fTB", b/TB)
    23  	case b >= GB:
    24  		return fmt.Sprintf("%.2fGB", b/GB)
    25  	case b >= MB:
    26  		return fmt.Sprintf("%.1fMB", b/MB)
    27  	case b >= KB:
    28  		return fmt.Sprintf("%.0fKB", b/KB)
    29  	}
    30  	return fmt.Sprintf("%.0fB", b)
    31  }
    32  
    33  type Transfer interface {
    34  	Transferred() ByteSize
    35  	Length() ByteSize
    36  }
    37  
    38  // ReaderTransfer monitors how much of a reader has been read
    39  type ReaderTransfer struct {
    40  	length ByteSize
    41  	read   uint64
    42  	reader io.Reader
    43  }
    44  
    45  // NewReaderTransfer instantiates a ReadTransfer struct
    46  func NewReaderTransfer(reader io.Reader, length int) *ReaderTransfer {
    47  	rt := new(ReaderTransfer)
    48  	rt.length = ByteSize(length)
    49  	rt.read = 0
    50  	rt.reader = reader
    51  	return rt
    52  }
    53  
    54  func (rt *ReaderTransfer) Read(p []byte) (int, error) {
    55  	n, err := rt.reader.Read(p)
    56  	atomic.AddUint64(&rt.read, uint64(n))
    57  	return n, err
    58  }
    59  
    60  func (rt *ReaderTransfer) Transferred() ByteSize {
    61  	return ByteSize(atomic.LoadUint64(&rt.read))
    62  }
    63  
    64  func (rt *ReaderTransfer) Length() ByteSize {
    65  	return rt.length
    66  }
    67  
    68  // WriteCloserTransfer monitors how much of a WriteCloser has been written
    69  type WriteCloserTransfer struct {
    70  	length      ByteSize
    71  	written     uint64
    72  	writeCloser io.WriteCloser
    73  }
    74  
    75  // NewWriteCloserTransfer instantiates a new WriteCloserTransfer
    76  func NewWriteCloserTransfer(writeCloser io.WriteCloser, length int) *WriteCloserTransfer {
    77  	wct := new(WriteCloserTransfer)
    78  	wct.length = ByteSize(length)
    79  	wct.written = 0
    80  	wct.writeCloser = writeCloser
    81  	return wct
    82  }
    83  
    84  func (wct *WriteCloserTransfer) Write(p []byte) (int, error) {
    85  	n, err := wct.writeCloser.Write(p)
    86  	atomic.AddUint64(&wct.written, uint64(n))
    87  	return n, err
    88  }
    89  
    90  func (wct *WriteCloserTransfer) Close() error {
    91  	return wct.writeCloser.Close()
    92  }
    93  
    94  func (wct *WriteCloserTransfer) Transferred() ByteSize {
    95  	return ByteSize(atomic.LoadUint64(&wct.written))
    96  }
    97  
    98  func (wct *WriteCloserTransfer) Length() ByteSize {
    99  	return wct.length
   100  }