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         "strings"
  25         "sync"
  26         "time"
  27
  28         "github.com/dustin/go-humanize"
  29         "go.cypherpunks.ru/nncp/v5/uilive"
  30 )
  31
  32 func init() {
  33         uilive.Out = os.Stderr
  34 }
  35
  36 var progressBars = make(map[string]*ProgressBar)
  37 var progressBarsLock sync.RWMutex
  38
  39 type ProgressBar struct {
  40         w       *uilive.Writer
  41         hash    string
  42         started time.Time
  43         initial int64
  44         full    int64
  45 }
  46
  47 func ProgressBarNew(initial, full int64) *ProgressBar {
  48         pb := ProgressBar{
  49                 w:       uilive.New(),
  50                 started: time.Now(),
  51                 initial: initial,
  52                 full:    full,
  53         }
  54         pb.w.Start()
  55         return &pb
  56 }
  57
  58 func (pb ProgressBar) Render(what string, size int64) {
  59         now := time.Now().UTC()
  60         timeDiff := now.Sub(pb.started).Seconds()
  61         if timeDiff == 0 {
  62                 timeDiff = 1
  63         }
  64         percentage := int64(100)
  65         if pb.full > 0 {
  66                 percentage = 100 * size / pb.full
  67         }
  68         fmt.Fprintf(
  69                 pb.w, "%s %s %s/%s %d%% (%s/sec)\n",
  70                 now.Format(time.RFC3339), what,
  71                 humanize.IBytes(uint64(size)),
  72                 humanize.IBytes(uint64(pb.full)),
  73                 percentage,
  74                 humanize.IBytes(uint64(float64(size-pb.initial)/timeDiff)),
  75         )
  76 }
  77
  78 func (pb ProgressBar) Kill() {
  79         pb.w.Stop()
  80 }
  81
  82 func CopyProgressed(
  83         dst io.Writer,
  84         src io.Reader,
  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(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(sds SDS) {
 122         pkt := sds["pkt"].(string)
 123         var size int64
 124         if sizeI, exists := sds["size"]; exists {
 125                 size = sizeI.(int64)
 126         }
 127         fullsize := sds["fullsize"].(int64)
 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         if xx, exists := sds["xx"]; exists {
 142                 what = strings.Title(xx.(string)) + " " + what
 143         }
 144         pb.Render(what, size)
 145         if size >= fullsize {
 146                 pb.Kill()
 147                 progressBarsLock.Lock()
 148                 delete(progressBars, pkt)
 149                 progressBarsLock.Unlock()
 150         }
 151 }