0
0
Fork 0
mirror of https://github.com/schollz/croc.git synced 2025-10-11 13:21:00 +02:00

Use a home-rolled progressbar

This commit is contained in:
Zack Scholl 2017-10-26 12:23:50 -06:00
parent 6fd3a40700
commit e3ff37cce8

View file

@ -14,9 +14,9 @@ import (
"time" "time"
humanize "github.com/dustin/go-humanize" humanize "github.com/dustin/go-humanize"
"github.com/schollz/progressbar"
"github.com/verybluebot/tarinator-go" "github.com/verybluebot/tarinator-go"
"github.com/gosuri/uiprogress"
"github.com/pkg/errors" "github.com/pkg/errors"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@ -33,7 +33,7 @@ type Connection struct {
Debug bool Debug bool
DontEncrypt bool DontEncrypt bool
Wait bool Wait bool
bars []*uiprogress.Bar bar *progressbar.ProgressBar
rate int rate int
} }
@ -218,9 +218,8 @@ func (c *Connection) runClient() error {
var wg sync.WaitGroup var wg sync.WaitGroup
wg.Add(c.NumberOfConnections) wg.Add(c.NumberOfConnections)
uiprogress.Start()
if !c.Debug { if !c.Debug {
c.bars = make([]*uiprogress.Bar, c.NumberOfConnections) c.bar = progressbar.New(c.File.Size)
} }
type responsesStruct struct { type responsesStruct struct {
gotTimeout bool gotTimeout bool
@ -294,6 +293,7 @@ func (c *Connection) runClient() error {
// Write data from file // Write data from file
logger.Debug("send file") logger.Debug("send file")
startTime = time.Now() startTime = time.Now()
c.bar.Reset()
if err := c.sendFile(id, connection); err != nil { if err := c.sendFile(id, connection); err != nil {
log.Error(err) log.Error(err)
} }
@ -388,6 +388,8 @@ func (c *Connection) runClient() error {
fmt.Printf("\n\nReceiving (<-%s)..\n", sendersAddress) fmt.Printf("\n\nReceiving (<-%s)..\n", sendersAddress)
} }
startTime = time.Now() startTime = time.Now()
c.bar.Set(c.File.Size)
c.bar.Reset()
if err := c.receiveFile(id, connection); err != nil { if err := c.receiveFile(id, connection); err != nil {
log.Error(errors.Wrap(err, "Problem receiving the file: ")) log.Error(errors.Wrap(err, "Problem receiving the file: "))
} }
@ -452,7 +454,7 @@ func (c *Connection) runClient() error {
return fmt.Errorf("\nUh oh! %s is corrupted! Sorry, try again.\n", c.File.Name) return fmt.Errorf("\nUh oh! %s is corrupted! Sorry, try again.\n", c.File.Name)
} }
if c.File.IsDir { // if the file was originally a dir if c.File.IsDir { // if the file was originally a dir
fmt.Print("decompressing folder") fmt.Print("\ndecompressing folder")
log.Debug("untarring " + c.File.Name) log.Debug("untarring " + c.File.Name)
err := tarinator.UnTarinate(c.Path, path.Join(c.Path, c.File.Name)) err := tarinator.UnTarinate(c.Path, path.Join(c.Path, c.File.Name))
@ -518,18 +520,10 @@ func (c *Connection) receiveFile(id int, connection net.Conn) error {
} }
defer newFile.Close() defer newFile.Close()
if !c.Debug {
c.bars[id] = uiprogress.AddBar(int(chunkSize)/1024 + 1).AppendCompleted().PrependElapsed()
c.bars[id].Width = 40
}
logger.Debug("waiting for file") logger.Debug("waiting for file")
var receivedBytes int64 var receivedBytes int64
receivedFirstBytes := false receivedFirstBytes := false
for { for {
if !c.Debug {
c.bars[id].Incr()
}
if (chunkSize - receivedBytes) < BUFFERSIZE { if (chunkSize - receivedBytes) < BUFFERSIZE {
logger.Debug("at the end") logger.Debug("at the end")
io.CopyN(newFile, connection, (chunkSize - receivedBytes)) io.CopyN(newFile, connection, (chunkSize - receivedBytes))
@ -538,6 +532,9 @@ func (c *Connection) receiveFile(id int, connection net.Conn) error {
logger.Debug("empty remaining bytes from network buffer") logger.Debug("empty remaining bytes from network buffer")
connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-chunkSize)) connection.Read(make([]byte, (receivedBytes+BUFFERSIZE)-chunkSize))
} }
if !c.Debug {
c.bar.Add(int((chunkSize - receivedBytes)))
}
break break
} }
io.CopyN(newFile, connection, BUFFERSIZE) io.CopyN(newFile, connection, BUFFERSIZE)
@ -546,6 +543,9 @@ func (c *Connection) receiveFile(id int, connection net.Conn) error {
receivedFirstBytes = true receivedFirstBytes = true
logger.Debug("Receieved first bytes!") logger.Debug("Receieved first bytes!")
} }
if !c.Debug {
c.bar.Add(BUFFERSIZE)
}
} }
logger.Debug("received file") logger.Debug("received file")
return nil return nil
@ -576,13 +576,6 @@ func (c *Connection) sendFile(id int, connection net.Conn) error {
return errors.Wrap(err, "Problem sending chunk data: ") return errors.Wrap(err, "Problem sending chunk data: ")
} }
// show the progress
if !c.Debug {
logger.Debug("going to show progress")
c.bars[id] = uiprogress.AddBar(int(fi.Size())).AppendCompleted().PrependElapsed()
c.bars[id].Width = 40
}
// rate limit the bandwidth // rate limit the bandwidth
logger.Debug("determining rate limiting") logger.Debug("determining rate limiting")
bufferSizeInKilobytes := BUFFERSIZE / 1024 bufferSizeInKilobytes := BUFFERSIZE / 1024
@ -598,7 +591,7 @@ func (c *Connection) sendFile(id int, connection net.Conn) error {
connection.Write(sendBuffer) connection.Write(sendBuffer)
totalBytesSent += n totalBytesSent += n
if !c.Debug { if !c.Debug {
c.bars[id].Set(totalBytesSent) c.bar.Add(n)
} }
if err == io.EOF { if err == io.EOF {
//End of file reached, break out of for loop //End of file reached, break out of for loop