github.com/aliyun/aliyun-oss-go-sdk@v3.0.2+incompatible/oss/progress.go (about)

     1  package oss
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  // ProgressEventType defines transfer progress event type
     8  type ProgressEventType int
     9  
    10  const (
    11  	// TransferStartedEvent transfer started, set TotalBytes
    12  	TransferStartedEvent ProgressEventType = 1 + iota
    13  	// TransferDataEvent transfer data, set ConsumedBytes and TotalBytes
    14  	TransferDataEvent
    15  	// TransferCompletedEvent transfer completed
    16  	TransferCompletedEvent
    17  	// TransferFailedEvent transfer encounters an error
    18  	TransferFailedEvent
    19  )
    20  
    21  // ProgressEvent defines progress event
    22  type ProgressEvent struct {
    23  	ConsumedBytes int64
    24  	TotalBytes    int64
    25  	RwBytes       int64
    26  	EventType     ProgressEventType
    27  }
    28  
    29  // ProgressListener listens progress change
    30  type ProgressListener interface {
    31  	ProgressChanged(event *ProgressEvent)
    32  }
    33  
    34  // -------------------- Private --------------------
    35  
    36  func newProgressEvent(eventType ProgressEventType, consumed, total int64, rwBytes int64) *ProgressEvent {
    37  	return &ProgressEvent{
    38  		ConsumedBytes: consumed,
    39  		TotalBytes:    total,
    40  		RwBytes:       rwBytes,
    41  		EventType:     eventType}
    42  }
    43  
    44  // publishProgress
    45  func publishProgress(listener ProgressListener, event *ProgressEvent) {
    46  	if listener != nil && event != nil {
    47  		listener.ProgressChanged(event)
    48  	}
    49  }
    50  
    51  type readerTracker struct {
    52  	completedBytes int64
    53  }
    54  
    55  type teeReader struct {
    56  	reader        io.Reader
    57  	writer        io.Writer
    58  	listener      ProgressListener
    59  	consumedBytes int64
    60  	totalBytes    int64
    61  	tracker       *readerTracker
    62  }
    63  
    64  // TeeReader returns a Reader that writes to w what it reads from r.
    65  // All reads from r performed through it are matched with
    66  // corresponding writes to w.  There is no internal buffering -
    67  // the write must complete before the read completes.
    68  // Any error encountered while writing is reported as a read error.
    69  func TeeReader(reader io.Reader, writer io.Writer, totalBytes int64, listener ProgressListener, tracker *readerTracker) io.ReadCloser {
    70  	return &teeReader{
    71  		reader:        reader,
    72  		writer:        writer,
    73  		listener:      listener,
    74  		consumedBytes: 0,
    75  		totalBytes:    totalBytes,
    76  		tracker:       tracker,
    77  	}
    78  }
    79  
    80  func (t *teeReader) Read(p []byte) (n int, err error) {
    81  	n, err = t.reader.Read(p)
    82  
    83  	// Read encountered error
    84  	if err != nil && err != io.EOF {
    85  		event := newProgressEvent(TransferFailedEvent, t.consumedBytes, t.totalBytes, 0)
    86  		publishProgress(t.listener, event)
    87  	}
    88  
    89  	if n > 0 {
    90  		t.consumedBytes += int64(n)
    91  		// CRC
    92  		if t.writer != nil {
    93  			if n, err := t.writer.Write(p[:n]); err != nil {
    94  				return n, err
    95  			}
    96  		}
    97  		// Progress
    98  		if t.listener != nil {
    99  			event := newProgressEvent(TransferDataEvent, t.consumedBytes, t.totalBytes, int64(n))
   100  			publishProgress(t.listener, event)
   101  		}
   102  		// Track
   103  		if t.tracker != nil {
   104  			t.tracker.completedBytes = t.consumedBytes
   105  		}
   106  	}
   107  
   108  	return
   109  }
   110  
   111  func (t *teeReader) Close() error {
   112  	if rc, ok := t.reader.(io.ReadCloser); ok {
   113  		return rc.Close()
   114  	}
   115  	return nil
   116  }