src/progress.go

   1 /*
   2 NNCP -- Node to Node copy, utilities for store-and-forward data exchange
   3 Copyright (C) 2016-2021 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/v7/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         les LEs,
  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                                         Progress(prgrsPrefix, append(les, LE{"Size", written}))
  99                                 }
 100                         }
 101                         if ew != nil {
 102                                 err = ew
 103                                 break
 104                         }
 105                         if nr != nw {
 106                                 err = io.ErrShortWrite
 107                                 break
 108                         }
 109                 }
 110                 if er != nil {
 111                         if er != io.EOF {
 112                                 err = er
 113                         }
 114                         break
 115                 }
 116         }
 117         return
 118 }
 119
 120 func Progress(prefix string, les LEs) {
 121         var size int64
 122         var fullsize int64
 123         var pkt string
 124         for _, le := range les {
 125                 switch le.K {
 126                 case "Size":
 127                         size = le.V.(int64)
 128                 case "FullSize":
 129                         fullsize = le.V.(int64)
 130                 case "Pkt":
 131                         pkt = le.V.(string)
 132                 }
 133         }
 134         progressBarsLock.RLock()
 135         pb := progressBars[pkt]
 136         progressBarsLock.RUnlock()
 137         if pb == nil {
 138                 progressBarsLock.Lock()
 139                 pb = ProgressBarNew(size, fullsize)
 140                 progressBars[pkt] = pb
 141                 progressBarsLock.Unlock()
 142         }
 143         what := pkt
 144         if len(what) >= Base32Encoded32Len { // Base32 encoded
 145                 what = what[:16] + ".." + what[len(what)-16:]
 146         }
 147         what = prefix + " " + what
 148         pb.Render(what, size)
 149         if size >= fullsize {
 150                 pb.Kill()
 151                 progressBarsLock.Lock()
 152                 delete(progressBars, pkt)
 153                 progressBarsLock.Unlock()
 154         }
 155 }
 156
 157 func ProgressKill(pkt string) {
 158         progressBarsLock.Lock()
 159         pb := progressBars[pkt]
 160         if pb != nil {
 161                 pb.Kill()
 162                 delete(progressBars, pkt)
 163         }
 164         progressBarsLock.Unlock()
 165 }