github.com/cavaliergopher/grab/v3@v3.0.1/pkg/grabui/console_client.go (about)

     1  package grabui
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/cavaliergopher/grab/v3"
    11  )
    12  
    13  type ConsoleClient struct {
    14  	mu                            sync.Mutex
    15  	client                        *grab.Client
    16  	succeeded, failed, inProgress int
    17  	responses                     []*grab.Response
    18  }
    19  
    20  func NewConsoleClient(client *grab.Client) *ConsoleClient {
    21  	return &ConsoleClient{
    22  		client: client,
    23  	}
    24  }
    25  
    26  func (c *ConsoleClient) Do(
    27  	ctx context.Context,
    28  	workers int,
    29  	reqs ...*grab.Request,
    30  ) <-chan *grab.Response {
    31  	// buffer size prevents slow receivers causing back pressure
    32  	pump := make(chan *grab.Response, len(reqs))
    33  
    34  	go func() {
    35  		c.mu.Lock()
    36  		defer c.mu.Unlock()
    37  
    38  		c.failed = 0
    39  		c.inProgress = 0
    40  		c.succeeded = 0
    41  		c.responses = make([]*grab.Response, 0, len(reqs))
    42  		if c.client == nil {
    43  			c.client = grab.DefaultClient
    44  		}
    45  
    46  		fmt.Printf("Downloading %d files...\n", len(reqs))
    47  		respch := c.client.DoBatch(workers, reqs...)
    48  		t := time.NewTicker(200 * time.Millisecond)
    49  		defer t.Stop()
    50  
    51  	Loop:
    52  		for {
    53  			select {
    54  			case <-ctx.Done():
    55  				break Loop
    56  
    57  			case resp := <-respch:
    58  				if resp != nil {
    59  					// a new response has been received and has started downloading
    60  					c.responses = append(c.responses, resp)
    61  					pump <- resp // send to caller
    62  				} else {
    63  					// channel is closed - all downloads are complete
    64  					break Loop
    65  				}
    66  
    67  			case <-t.C:
    68  				// update UI on clock tick
    69  				c.refresh()
    70  			}
    71  		}
    72  
    73  		c.refresh()
    74  		close(pump)
    75  
    76  		fmt.Printf(
    77  			"Finished %d successful, %d failed, %d incomplete.\n",
    78  			c.succeeded,
    79  			c.failed,
    80  			c.inProgress)
    81  	}()
    82  	return pump
    83  }
    84  
    85  // refresh prints the progress of all downloads to the terminal
    86  func (c *ConsoleClient) refresh() {
    87  	// clear lines for incomplete downloads
    88  	if c.inProgress > 0 {
    89  		fmt.Printf("\033[%dA\033[K", c.inProgress)
    90  	}
    91  
    92  	// print newly completed downloads
    93  	for i, resp := range c.responses {
    94  		if resp != nil && resp.IsComplete() {
    95  			if resp.Err() != nil {
    96  				c.failed++
    97  				fmt.Fprintf(os.Stderr, "Error downloading %s: %v\n",
    98  					resp.Request.URL(),
    99  					resp.Err())
   100  			} else {
   101  				c.succeeded++
   102  				fmt.Printf("Finished %s %s / %s (%d%%)\n",
   103  					resp.Filename,
   104  					byteString(resp.BytesComplete()),
   105  					byteString(resp.Size()),
   106  					int(100*resp.Progress()))
   107  			}
   108  			c.responses[i] = nil
   109  		}
   110  	}
   111  
   112  	// print progress for incomplete downloads
   113  	c.inProgress = 0
   114  	for _, resp := range c.responses {
   115  		if resp != nil {
   116  			fmt.Printf("Downloading %s %s / %s (%d%%) - %s ETA: %s \033[K\n",
   117  				resp.Filename,
   118  				byteString(resp.BytesComplete()),
   119  				byteString(resp.Size()),
   120  				int(100*resp.Progress()),
   121  				bpsString(resp.BytesPerSecond()),
   122  				etaString(resp.ETA()))
   123  			c.inProgress++
   124  		}
   125  	}
   126  }
   127  
   128  func bpsString(n float64) string {
   129  	if n < 1e3 {
   130  		return fmt.Sprintf("%.02fBps", n)
   131  	}
   132  	if n < 1e6 {
   133  		return fmt.Sprintf("%.02fKB/s", n/1e3)
   134  	}
   135  	if n < 1e9 {
   136  		return fmt.Sprintf("%.02fMB/s", n/1e6)
   137  	}
   138  	return fmt.Sprintf("%.02fGB/s", n/1e9)
   139  }
   140  
   141  func byteString(n int64) string {
   142  	if n < 1<<10 {
   143  		return fmt.Sprintf("%dB", n)
   144  	}
   145  	if n < 1<<20 {
   146  		return fmt.Sprintf("%dKB", n>>10)
   147  	}
   148  	if n < 1<<30 {
   149  		return fmt.Sprintf("%dMB", n>>20)
   150  	}
   151  	if n < 1<<40 {
   152  		return fmt.Sprintf("%dGB", n>>30)
   153  	}
   154  	return fmt.Sprintf("%dTB", n>>40)
   155  }
   156  
   157  func etaString(eta time.Time) string {
   158  	d := eta.Sub(time.Now())
   159  	if d < time.Second {
   160  		return "<1s"
   161  	}
   162  	// truncate to 1s resolution
   163  	d /= time.Second
   164  	d *= time.Second
   165  	return d.String()
   166  }