From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 21 May 2021 01:49:57 +0000 (+1000)
Subject: Track chunks read for webseed peers too
X-Git-Tag: v1.29.0~31^2~9
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=c895a21a165090d8a37ac43350e46eb67a00dd4d;p=btrtrc.git

Track chunks read for webseed peers too
---

diff --git a/conn_stats.go b/conn_stats.go
index d2d52003..0c5bfc78 100644
--- a/conn_stats.go
+++ b/conn_stats.go
@@ -79,14 +79,9 @@ func (cs *ConnStats) wroteMsg(msg *pp.Message) {
 	}
 }
 
-func (cs *ConnStats) readMsg(msg *pp.Message) {
-	// We want to also handle extended metadata pieces here, but we wouldn't
-	// have decoded the extended payload yet.
-	switch msg.Type {
-	case pp.Piece:
-		cs.ChunksRead.Add(1)
-		cs.BytesReadData.Add(int64(len(msg.Piece)))
-	}
+func (cs *ConnStats) receivedChunk(size int64) {
+	cs.ChunksRead.Add(1)
+	cs.BytesReadData.Add(size)
 }
 
 func (cs *ConnStats) incrementPiecesDirtiedGood() {
diff --git a/peerconn.go b/peerconn.go
index 40848e37..46756320 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -856,10 +856,6 @@ func (cn *PeerConn) wroteMsg(msg *pp.Message) {
 	cn.allStats(func(cs *ConnStats) { cs.wroteMsg(msg) })
 }
 
-func (cn *PeerConn) readMsg(msg *pp.Message) {
-	cn.allStats(func(cs *ConnStats) { cs.readMsg(msg) })
-}
-
 // After handshake, we know what Torrent and Client stats to include for a
 // connection.
 func (cn *Peer) postHandshakeStats(f func(*ConnStats)) {
@@ -1065,7 +1061,6 @@ func (c *PeerConn) mainReadLoop() (err error) {
 		if err != nil {
 			return err
 		}
-		c.readMsg(&msg)
 		c.lastMessageReceived = time.Now()
 		if msg.Keepalive {
 			receivedKeepalives.Add(1)
@@ -1105,6 +1100,7 @@ func (c *PeerConn) mainReadLoop() (err error) {
 			r := newRequestFromMessage(&msg)
 			err = c.onReadRequest(r)
 		case pp.Piece:
+			c.doChunkReadStats(int64(len(msg.Piece)))
 			err = c.receiveChunk(&msg)
 			if len(msg.Piece) == int(t.chunkSize) {
 				t.chunkPool.Put(&msg.Piece)
@@ -1253,10 +1249,12 @@ func (cn *PeerConn) rw() io.ReadWriter {
 	}{cn.r, cn.w}
 }
 
+func (c *Peer) doChunkReadStats(size int64) {
+	c.allStats(func(cs *ConnStats) { cs.receivedChunk(size) })
+}
+
 // Handle a received chunk from a peer.
 func (c *Peer) receiveChunk(msg *pp.Message) error {
-	t := c.t
-	cl := t.cl
 	chunksReceived.Add("total", 1)
 
 	req := newRequestFromMessage(msg)
@@ -1296,6 +1294,9 @@ func (c *Peer) receiveChunk(msg *pp.Message) error {
 		}
 	}
 
+	t := c.t
+	cl := t.cl
+
 	// Do we actually want this chunk?
 	if t.haveChunk(req) {
 		chunksReceived.Add("wasted", 1)
diff --git a/webseed-peer.go b/webseed-peer.go
index 7af892f0..f2ccbff5 100644
--- a/webseed-peer.go
+++ b/webseed-peer.go
@@ -110,6 +110,10 @@ func (ws *webseedPeer) onClose() {
 
 func (ws *webseedPeer) requestResultHandler(r Request, webseedRequest webseed.Request) {
 	result := <-webseedRequest.Result
+	// We do this here rather than inside receiveChunk, since we want to count errors too. I'm not
+	// sure if we can divine which errors indicate cancellation on our end without hitting the
+	// network though.
+	ws.peer.doChunkReadStats(int64(len(result.Bytes)))
 	ws.peer.t.cl.lock()
 	defer ws.peer.t.cl.unlock()
 	if result.Err != nil {