src/progress.go

   1 /*
   2 NNCP -- Node to Node copy, utilities for store-and-forward data exchange
   3 Copyright (C) 2016-2019 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package nncp
  19
  20 import (
  21         "fmt"
  22         "io"
  23         "os"
  24         "sync"
  25         "time"
  26
  27         "github.com/dustin/go-humanize"
  28         "go.cypherpunks.ru/nncp/v5/uilive"
  29 )
  30
  31 func init() {
  32         uilive.Out = os.Stderr
  33 }
  34
  35 var progressBars = make(map[string]*ProgressBar)
  36 var progressBarsLock sync.RWMutex
  37
  38 type ProgressBar struct {
  39         w       *uilive.Writer
  40         hash    string
  41         started time.Time
  42         initial int64
  43         full    int64
  44 }
  45
  46 func ProgressBarNew(initial, full int64) *ProgressBar {
  47         pb := ProgressBar{
  48                 w:       uilive.New(),
  49                 started: time.Now(),
  50                 initial: initial,
  51                 full:    full,
  52         }
  53         pb.w.Start()
  54         return &pb
  55 }
  56
  57 func (pb ProgressBar) Render(what string, size int64) {
  58         now := time.Now().UTC()
  59         timeDiff := now.Sub(pb.started).Seconds()
  60         if timeDiff == 0 {
  61                 timeDiff = 1
  62         }
  63         percentage := int64(100)
  64         if pb.full > 0 {
  65                 percentage = 100 * size / pb.full
  66         }
  67         fmt.Fprintf(
  68                 pb.w, "%s %s %s/%s %d%% (%s/sec)\n",
  69                 now.Format(time.RFC3339), what,
  70                 humanize.IBytes(uint64(size)),
  71                 humanize.IBytes(uint64(pb.full)),
  72                 percentage,
  73                 humanize.IBytes(uint64(float64(size-pb.initial)/timeDiff)),
  74         )
  75 }
  76
  77 func (pb ProgressBar) Kill() {
  78         pb.w.Stop()
  79 }
  80
  81 func CopyProgressed(
  82         dst io.Writer,
  83         src io.Reader,
  84         prgrsPrefix string,
  85         sds SDS,
  86         showPrgrs bool,
  87 ) (written int64, err error) {
  88         buf := make([]byte, EncBlkSize)
  89         var nr, nw int
  90         var er, ew error
  91         for {
  92                 nr, er = src.Read(buf)
  93                 if nr > 0 {
  94                         nw, ew = dst.Write(buf[:nr])
  95                         if nw > 0 {
  96                                 written += int64(nw)
  97                                 if showPrgrs {
  98                                         sds["size"] = written
  99                                         Progress(prgrsPrefix, sds)
 100                                 }
 101                         }
 102                         if ew != nil {
 103                                 err = ew
 104                                 break
 105                         }
 106                         if nr != nw {
 107                                 err = io.ErrShortWrite
 108                                 break
 109                         }
 110                 }
 111                 if er != nil {
 112                         if er != io.EOF {
 113                                 err = er
 114                         }
 115                         break
 116                 }
 117         }
 118         return
 119 }
 120
 121 func Progress(prefix string, sds SDS) {
 122         var size int64
 123         if sizeI, exists := sds["size"]; exists {
 124                 size = sizeI.(int64)
 125         }
 126         fullsize := sds["fullsize"].(int64)
 127         pkt := sds["pkt"].(string)
 128         progressBarsLock.RLock()
 129         pb, exists := progressBars[pkt]
 130         progressBarsLock.RUnlock()
 131         if !exists {
 132                 progressBarsLock.Lock()
 133                 pb = ProgressBarNew(size, fullsize)
 134                 progressBars[pkt] = pb
 135                 progressBarsLock.Unlock()
 136         }
 137         what := pkt
 138         if len(what) >= 52 { // Base32 encoded
 139                 what = what[:16] + ".." + what[len(what)-16:]
 140         }
 141         what = prefix + " " + what
 142         pb.Render(what, size)
 143         if size >= fullsize {
 144                 pb.Kill()
 145                 progressBarsLock.Lock()
 146                 delete(progressBars, pkt)
 147                 progressBarsLock.Unlock()
 148         }
 149 }