From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 5 Jul 2016 05:52:33 +0000 (+1000)
Subject: Add counters for uploads
X-Git-Tag: v1.0.0~669
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=12191dbfa36559e375729926c38915b101c9b302;p=btrtrc.git

Add counters for uploads

Fixes #71.
---

diff --git a/client_test.go b/client_test.go
index 6cefd18e..648b55c2 100644
--- a/client_test.go
+++ b/client_test.go
@@ -306,6 +306,8 @@ type testClientTransferParams struct {
 	SeederStorage                       func(string) storage.Client
 }
 
+// Creates a seeder and a leecher, and ensures the data transfers when a read
+// is attempted on the leecher.
 func testClientTransfer(t *testing.T, ps testClientTransferParams) {
 	greetingTempDir, mi := testutil.GreetingTestTorrent()
 	defer os.RemoveAll(greetingTempDir)
@@ -322,7 +324,7 @@ func testClientTransfer(t *testing.T, ps testClientTransferParams) {
 	if ps.ExportClientStatus {
 		testutil.ExportStatusWriter(seeder, "s")
 	}
-	_, new, err := seeder.AddTorrentSpec(TorrentSpecFromMetaInfo(mi))
+	seederTorrent, new, err := seeder.AddTorrentSpec(TorrentSpecFromMetaInfo(mi))
 	require.NoError(t, err)
 	assert.True(t, new)
 	leecherDataDir, err := ioutil.TempDir("", "")
@@ -362,14 +364,24 @@ func testClientTransfer(t *testing.T, ps testClientTransferParams) {
 	if ps.SetReadahead {
 		r.SetReadahead(ps.Readahead)
 	}
-	for range iter.N(2) {
-		pos, err := r.Seek(0, os.SEEK_SET)
-		assert.NoError(t, err)
-		assert.EqualValues(t, 0, pos)
-		_greeting, err := ioutil.ReadAll(r)
-		assert.NoError(t, err)
-		assert.EqualValues(t, testutil.GreetingFileContents, _greeting)
-	}
+	assertReadAllGreeting(t, r)
+	// After one read through, we can assume certain torrent statistics.
+	assert.EqualValues(t, 13, seederTorrent.Stats().DataBytesSent)
+	assert.EqualValues(t, 8, seederTorrent.Stats().ChunksSent)
+	// This is not a strict requirement. It is however interesting to follow.
+	assert.EqualValues(t, 261, seederTorrent.Stats().BytesSent)
+	// Read through again for the cases where the torrent data size exceed the
+	// size of the cache.
+	assertReadAllGreeting(t, r)
+}
+
+func assertReadAllGreeting(t *testing.T, r io.ReadSeeker) {
+	pos, err := r.Seek(0, os.SEEK_SET)
+	assert.NoError(t, err)
+	assert.EqualValues(t, 0, pos)
+	_greeting, err := ioutil.ReadAll(r)
+	assert.NoError(t, err)
+	assert.EqualValues(t, testutil.GreetingFileContents, _greeting)
 }
 
 // Check that after completing leeching, a leecher transitions to a seeding
diff --git a/conn_stats.go b/conn_stats.go
new file mode 100644
index 00000000..5455b49a
--- /dev/null
+++ b/conn_stats.go
@@ -0,0 +1,23 @@
+package torrent
+
+import (
+	pp "github.com/anacrolix/torrent/peer_protocol"
+)
+
+type ConnStats struct {
+	ChunksSent    int64 // Num piece messages sent.
+	BytesSent     int64 // Total bytes sent.
+	DataBytesSent int64 // Data-only bytes sent.
+}
+
+func (cs *ConnStats) wroteMsg(msg pp.Message) {
+	switch msg.Type {
+	case pp.Piece:
+		cs.ChunksSent++
+		cs.DataBytesSent += int64(len(msg.Piece))
+	}
+}
+
+func (cs *ConnStats) wroteBytes(b []byte) {
+	cs.BytesSent += int64(len(b))
+}
diff --git a/connection.go b/connection.go
index 3f0c4756..7f8e10b4 100644
--- a/connection.go
+++ b/connection.go
@@ -44,6 +44,7 @@ type connection struct {
 	uTP       bool
 	closed    missinggo.Event
 
+	stats                  ConnStats
 	UnwantedChunksReceived int
 	UsefulChunksReceived   int
 	chunksSent             int
@@ -425,6 +426,8 @@ func (cn *connection) writer(keepAliveTimeout time.Duration) {
 				panic("short write")
 			}
 			cn.mu().Lock()
+			cn.wroteMsg(msg)
+			cn.wroteBytes(b)
 		}
 		cn.outgoingUnbufferedMessagesNotEmpty.Clear()
 		cn.mu().Unlock()
@@ -640,3 +643,13 @@ func (c *connection) requestPendingMetadata() {
 		c.requestMetadataPiece(pending[i])
 	}
 }
+
+func (cn *connection) wroteMsg(msg pp.Message) {
+	cn.stats.wroteMsg(msg)
+	cn.t.stats.wroteMsg(msg)
+}
+
+func (cn *connection) wroteBytes(b []byte) {
+	cn.stats.wroteBytes(b)
+	cn.t.stats.wroteBytes(b)
+}
diff --git a/torrent.go b/torrent.go
index fa69665d..6af85be7 100644
--- a/torrent.go
+++ b/torrent.go
@@ -93,6 +93,7 @@ type Torrent struct {
 	completedPieces bitmap.Bitmap
 
 	connPieceInclinationPool sync.Pool
+	stats                    TorrentStats
 }
 
 func (t *Torrent) setDisplayName(dn string) {
@@ -1268,3 +1269,7 @@ func (t *Torrent) addPeers(peers []Peer) {
 		t.addPeer(p)
 	}
 }
+
+func (t *Torrent) Stats() TorrentStats {
+	return t.stats
+}
diff --git a/torrent_stats.go b/torrent_stats.go
new file mode 100644
index 00000000..3f34434a
--- /dev/null
+++ b/torrent_stats.go
@@ -0,0 +1,5 @@
+package torrent
+
+type TorrentStats struct {
+	ConnStats // Aggregates stats over all connections past and present.
+}