move reading to goroutine

This commit is contained in:
Zack Scholl 2018-09-25 09:55:35 -07:00
parent 75f7cdcf65
commit f62459e1a4
1 changed files with 52 additions and 33 deletions

View File

@ -49,6 +49,7 @@ func receive(forceSend int, serverAddress, serverTCP string, isLocal bool, c *we
var hash256 []byte var hash256 []byte
var otherIP string var otherIP string
var tcpConnection comm.Comm var tcpConnection comm.Comm
dataChan := make(chan []byte, 1024*1024)
useWebsockets := true useWebsockets := true
switch forceSend { switch forceSend {
@ -195,41 +196,21 @@ func receive(forceSend int, serverAddress, serverTCP string, isLocal bool, c *we
progressbar.OptionSetBytes(int(fstats.Size)), progressbar.OptionSetBytes(int(fstats.Size)),
progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSetWriter(os.Stderr),
) )
c.WriteMessage(websocket.BinaryMessage, []byte("ready")) finished := make(chan bool)
startTime := time.Now() go func(finished chan bool, dataChan chan []byte) (err error) {
var numBytes int
var bs []byte
for { for {
if useWebsockets { message := <-dataChan
var messageType int
// read from websockets
messageType, message, err = c.ReadMessage()
if messageType != websocket.BinaryMessage {
continue
}
} else {
// read from TCP connection
message, numBytes, bs, err = tcpConnection.Read()
// log.Debugf("message: %s", message)
}
if err != nil {
log.Error(err)
return err
}
if bytes.Equal(message, []byte("magic")) {
log.Debug("got magic")
break
}
// do decryption // do decryption
var enc crypt.Encryption var enc crypt.Encryption
err = json.Unmarshal(message, &enc) err = json.Unmarshal(message, &enc)
if err != nil { if err != nil {
log.Errorf("%s: [%s] [%+v] (%d/%d) %+v", err.Error(), message, message, len(message), numBytes, bs) // log.Errorf("%s: [%s] [%+v] (%d/%d) %+v", err.Error(), message, message, len(message), numBytes, bs)
log.Error(err)
return err return err
} }
decrypted, err := enc.Decrypt(sessionKey, !fstats.IsEncrypted) decrypted, err := enc.Decrypt(sessionKey, !fstats.IsEncrypted)
if err != nil { if err != nil {
log.Error(err)
return err return err
} }
@ -247,11 +228,49 @@ func receive(forceSend int, serverAddress, serverTCP string, isLocal bool, c *we
bytesWritten += n bytesWritten += n
// update the progress bar // update the progress bar
bar.Add(n) bar.Add(n)
if int64(bytesWritten) == fstats.Size {
// if int64(bytesWritten) == fstats.Size { log.Debug("finished")
// break break
// }
} }
}
finished <- true
return
}(finished, dataChan)
c.WriteMessage(websocket.BinaryMessage, []byte("ready"))
startTime := time.Now()
for {
if useWebsockets {
var messageType int
// read from websockets
messageType, message, err = c.ReadMessage()
if messageType != websocket.BinaryMessage {
continue
}
} else {
// read from TCP connection
message, _, _, err = tcpConnection.Read()
// log.Debugf("message: %s", message)
}
if err != nil {
log.Error(err)
return err
}
if bytes.Equal(message, []byte("magic")) {
log.Debug("got magic")
break
}
select {
case dataChan <- message:
continue
default:
log.Debug("blocked")
// no message sent
// block
dataChan <- message
}
}
_ = <-finished
c.WriteMessage(websocket.BinaryMessage, []byte("done")) c.WriteMessage(websocket.BinaryMessage, []byte("done"))
// we are finished // we are finished