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/v8/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         if showPrgrs {
 118                 for _, le := range les {
 119                         if le.K == "FullSize" {
 120                                 if le.V.(int64) == 0 {
 121                                         Progress(prgrsPrefix, append(
 122                                                 les, LE{"Size", written}, LE{"FullSize", written},
 123                                         ))
 124                                 }
 125                                 break
 126                         }
 127                 }
 128         }
 129         return
 130 }
 131
 132 func Progress(prefix string, les LEs) {
 133         var size int64
 134         var fullsize int64
 135         var pkt string
 136         for _, le := range les {
 137                 switch le.K {
 138                 case "Size":
 139                         size = le.V.(int64)
 140                 case "FullSize":
 141                         fullsize = le.V.(int64)
 142                 case "Pkt":
 143                         pkt = le.V.(string)
 144                 }
 145         }
 146         progressBarsLock.RLock()
 147         pb := progressBars[pkt]
 148         progressBarsLock.RUnlock()
 149         if pb == nil {
 150                 progressBarsLock.Lock()
 151                 pb = ProgressBarNew(size, fullsize)
 152                 progressBars[pkt] = pb
 153                 progressBarsLock.Unlock()
 154         }
 155         what := pkt
 156         if len(what) >= Base32Encoded32Len { // Base32 encoded
 157                 what = what[:16] + ".." + what[len(what)-16:]
 158         }
 159         what = prefix + " " + what
 160         pb.Render(what, size)
 161         if fullsize != 0 && size >= fullsize {
 162                 pb.Kill()
 163                 progressBarsLock.Lock()
 164                 delete(progressBars, pkt)
 165                 progressBarsLock.Unlock()
 166         }
 167 }
 168
 169 func ProgressKill(pkt string) {
 170         progressBarsLock.Lock()
 171         pb := progressBars[pkt]
 172         if pb != nil {
 173                 pb.Kill()
 174                 delete(progressBars, pkt)
 175         }
 176         progressBarsLock.Unlock()
 177 }