github.com/loggregator/cli@v6.33.1-0.20180224010324-82334f081791+incompatible/cf/net/progress_reader.go (about)

     1  package net
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"sync"
     7  	"time"
     8  
     9  	"code.cloudfoundry.org/cli/cf/formatters"
    10  	"code.cloudfoundry.org/cli/cf/terminal"
    11  )
    12  
    13  type ProgressReader struct {
    14  	ioReadSeeker   io.ReadSeeker
    15  	bytesRead      int64
    16  	total          int64
    17  	quit           chan bool
    18  	ui             terminal.UI
    19  	outputInterval time.Duration
    20  	mutex          sync.RWMutex
    21  }
    22  
    23  func NewProgressReader(readSeeker io.ReadSeeker, ui terminal.UI, outputInterval time.Duration) *ProgressReader {
    24  	return &ProgressReader{
    25  		ioReadSeeker:   readSeeker,
    26  		ui:             ui,
    27  		outputInterval: outputInterval,
    28  		mutex:          sync.RWMutex{},
    29  	}
    30  }
    31  
    32  func (progressReader *ProgressReader) Read(p []byte) (int, error) {
    33  	if progressReader.ioReadSeeker == nil {
    34  		return 0, os.ErrInvalid
    35  	}
    36  
    37  	n, err := progressReader.ioReadSeeker.Read(p)
    38  
    39  	if progressReader.total > int64(0) {
    40  		if n > 0 {
    41  			if progressReader.quit == nil {
    42  				progressReader.quit = make(chan bool)
    43  				go progressReader.printProgress(progressReader.quit)
    44  			}
    45  
    46  			progressReader.mutex.Lock()
    47  			progressReader.bytesRead += int64(n)
    48  			progressReader.mutex.Unlock()
    49  
    50  			if progressReader.total == progressReader.bytesRead {
    51  				progressReader.quit <- true
    52  				return n, err
    53  			}
    54  		}
    55  	}
    56  
    57  	return n, err
    58  }
    59  
    60  func (progressReader *ProgressReader) Seek(offset int64, whence int) (int64, error) {
    61  	return progressReader.ioReadSeeker.Seek(offset, whence)
    62  }
    63  
    64  func (progressReader *ProgressReader) printProgress(quit chan bool) {
    65  	timer := time.NewTicker(progressReader.outputInterval)
    66  
    67  	for {
    68  		select {
    69  		case <-quit:
    70  			//The spaces are there to ensure we overwrite the entire line
    71  			//before using the terminal printer to output Done Uploading
    72  			progressReader.ui.PrintCapturingNoOutput("\r                             ")
    73  			progressReader.ui.Say("\rDone uploading")
    74  			return
    75  		case <-timer.C:
    76  			progressReader.mutex.RLock()
    77  			progressReader.ui.PrintCapturingNoOutput("\r%s uploaded...", formatters.ByteSize(progressReader.bytesRead))
    78  			progressReader.mutex.RUnlock()
    79  		}
    80  	}
    81  }
    82  
    83  func (progressReader *ProgressReader) SetTotalSize(size int64) {
    84  	progressReader.total = size
    85  }