src/progress.go

   1 /*
   2 NNCP -- Node to Node copy, utilities for store-and-forward data exchange
   3 Copyright (C) 2016-2023 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         started time.Time
  41         initial int64
  42         full    int64
  43 }
  44
  45 func ProgressBarNew(initial, full int64) *ProgressBar {
  46         pb := ProgressBar{
  47                 w:       uilive.New(),
  48                 started: time.Now(),
  49                 initial: initial,
  50                 full:    full,
  51         }
  52         pb.w.Start()
  53         return &pb
  54 }
  55
  56 func (pb ProgressBar) Render(what string, size int64) {
  57         now := time.Now().UTC()
  58         timeDiff := now.Sub(pb.started).Seconds()
  59         if timeDiff == 0 {
  60                 timeDiff = 1
  61         }
  62         percentage := int64(100)
  63         if pb.full > 0 {
  64                 percentage = 100 * size / pb.full
  65         }
  66         fmt.Fprintf(
  67                 pb.w, "%s %s %s/%s %d%% (%s/sec)\n",
  68                 now.Format(time.RFC3339), what,
  69                 humanize.IBytes(uint64(size)),
  70                 humanize.IBytes(uint64(pb.full)),
  71                 percentage,
  72                 humanize.IBytes(uint64(float64(size-pb.initial)/timeDiff)),
  73         )
  74 }
  75
  76 func (pb ProgressBar) Kill() {
  77         pb.w.Stop()
  78 }
  79
  80 func CopyProgressed(
  81         dst io.Writer,
  82         src io.Reader,
  83         prgrsPrefix string,
  84         les LEs,
  85         showPrgrs bool,
  86 ) (written int64, err error) {
  87         buf := make([]byte, EncBlkSize)
  88         var nr, nw int
  89         var er, ew error
  90         for {
  91                 nr, er = src.Read(buf)
  92                 if nr > 0 {
  93                         nw, ew = dst.Write(buf[:nr])
  94                         if nw > 0 {
  95                                 written += int64(nw)
  96                                 if showPrgrs {
  97                                         Progress(prgrsPrefix, append(les, LE{"Size", written}))
  98                                 }
  99                         }
 100                         if ew != nil {
 101                                 err = ew
 102                                 break
 103                         }
 104                         if nr != nw {
 105                                 err = io.ErrShortWrite
 106                                 break
 107                         }
 108                 }
 109                 if er != nil {
 110                         if er != io.EOF {
 111                                 err = er
 112                         }
 113                         break
 114                 }
 115         }
 116         if showPrgrs {
 117                 for _, le := range les {
 118                         if le.K == "FullSize" {
 119                                 if le.V.(int64) == 0 {
 120                                         Progress(prgrsPrefix, append(
 121                                                 les, LE{"Size", written}, LE{"FullSize", written},
 122                                         ))
 123                                 }
 124                                 break
 125                         }
 126                 }
 127         }
 128         return
 129 }
 130
 131 func Progress(prefix string, les LEs) {
 132         var size int64
 133         var fullsize int64
 134         var pkt string
 135         for _, le := range les {
 136                 switch le.K {
 137                 case "Size":
 138                         size = le.V.(int64)
 139                 case "FullSize":
 140                         fullsize = le.V.(int64)
 141                 case "Pkt":
 142                         pkt = le.V.(string)
 143                 }
 144         }
 145         progressBarsLock.RLock()
 146         pb := progressBars[pkt]
 147         progressBarsLock.RUnlock()
 148         if pb == nil {
 149                 progressBarsLock.Lock()
 150                 pb = ProgressBarNew(size, fullsize)
 151                 progressBars[pkt] = pb
 152                 progressBarsLock.Unlock()
 153         }
 154         what := pkt
 155         if len(what) >= Base32Encoded32Len { // Base32 encoded
 156                 what = what[:16] + ".." + what[len(what)-16:]
 157         }
 158         what = prefix + " " + what
 159         pb.Render(what, size)
 160         if fullsize != 0 && size >= fullsize {
 161                 pb.Kill()
 162                 progressBarsLock.Lock()
 163                 delete(progressBars, pkt)
 164                 progressBarsLock.Unlock()
 165         }
 166 }
 167
 168 func ProgressKill(pkt string) {
 169         progressBarsLock.Lock()
 170         pb := progressBars[pkt]
 171         if pb != nil {
 172                 pb.Kill()
 173                 delete(progressBars, pkt)
 174         }
 175         progressBarsLock.Unlock()
 176 }