Merge pull request #36 from erigontech/upload_stats

Added peer bytes uploaded and rate

conn_stats.go

@@ -17,8 +17,9 @@ import (
 // Count, must be aligned on some platforms: See https://github.com/anacrolix/torrent/issues/262.
 type ConnStats struct {
 	// Total bytes on the wire. Includes handshakes and encryption.
-	BytesWritten     Count
-	BytesWrittenData Count
+	BytesWritten         Count
+	BytesWrittenData     Count
+	BytesWrittenDataRate Count //bytes per second
 
 	BytesRead                   Count
 	BytesReadData               Count
@@ -58,6 +59,10 @@ type Count struct {
 
 var _ fmt.Stringer = (*Count)(nil)
 
+func (me *Count) Store(n int64) {
+	atomic.StoreInt64(&me.n, n)
+}
+
 func (me *Count) Add(n int64) {
 	atomic.AddInt64(&me.n, n)
 }

peer.go

@@ -290,6 +290,18 @@ func (p *Peer) DownloadRate() float64 {
 	return p.downloadRate()
 }
 
+func (p *Peer) BytesUploaded() float64 {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return float64(p._stats.BytesWrittenDataRate.Int64())
+}
+
+func (p *Peer) UploadRate() int64 {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p._stats.BytesWrittenData.Int64()
+}
+
 func (cn *Peer) iterContiguousPieceRequests(f func(piece pieceIndex, count int), lockTorrent bool) {
 	var last Option[pieceIndex]
 	var count int

peerconn.go

@@ -59,6 +59,7 @@ type PeerConn struct {
 	PeerClientName   atomic.Value
 	uploadTimer      *time.Timer
 	pex              pexConnState
+	lastWriteTime    time.Time
 
 	// The pieces the peer has claimed to have.
 	_peerPieces roaring.Bitmap
@@ -633,7 +634,12 @@ func (cn *PeerConn) wroteMsg(msg *pp.Message) {
 			torrent.Add(fmt.Sprintf("Extended messages written for protocol %q", name), 1)
 		}
 	}
-	cn.allStats(func(cs *ConnStats) { cs.wroteMsg(msg) })
+	cn.allStats(func(cs *ConnStats) {
+		cs.wroteMsg(msg)
+		writeTime := time.Now()
+		cs.BytesWrittenDataRate.Store(int64(float64(len(msg.Piece)) / writeTime.Sub(cn.lastWriteTime).Seconds()))
+		cn.lastWriteTime = writeTime
+	})
 }
 
 func (cn *PeerConn) wroteBytes(n int64) {