src/progress.go

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