Add byte count tracking to event pairs

Amp-Thread-ID: https://ampcode.com/threads/T-019c8012-eaeb-768d-a264-5a704f3939ef
Co-authored-by: Amp <amp@ampcode.com>

3 files changed, 11 insertions, 4 deletions

diff --git a/internal/flamegraph/counter.go b/internal/flamegraph/counter.go
index 5a162c9..c12453a 100644
--- a/internal/flamegraph/counter.go
+++ b/internal/flamegraph/counter.go
@@ -8,7 +8,7 @@ type Counter struct {
 	Count          uint64
 	Duration       uint64
 	DurationToPrev uint64
-	Bytes          uint64 // TODO: implement
+	Bytes          uint64 // Bytes transferred (only set for read/write/transfer syscalls)
 }
 
 func (c Counter) add(other Counter) Counter {
diff --git a/internal/flamegraph/iordata.go b/internal/flamegraph/iordata.go
index 1c2b0f4..463ed48 100644
--- a/internal/flamegraph/iordata.go
+++ b/internal/flamegraph/iordata.go
@@ -51,7 +51,7 @@ func cloneString(s string) string {
 }
 
 func (iod iorData) addEventPair(ev *event.Pair) {
-	cnt := Counter{Count: 1, Duration: ev.Duration, DurationToPrev: ev.DurationToPrev}
+	cnt := Counter{Count: 1, Duration: ev.Duration, DurationToPrev: ev.DurationToPrev, Bytes: ev.Bytes}
 	iod.add(ev.FileName(), ev.EnterEv.GetTraceId(), strings.TrimSpace(ev.Comm), ev.EnterEv.GetPid(),
 		ev.EnterEv.GetTid(), ev.Flags(), cnt)
 }
diff --git a/internal/flamegraph/iordata_test.go b/internal/flamegraph/iordata_test.go
index f01828f..9957f9e 100644
--- a/internal/flamegraph/iordata_test.go
+++ b/internal/flamegraph/iordata_test.go
@@ -14,14 +14,14 @@ func TestAddPath(t *testing.T) {
 	pid := pidType(1234)
 	tid := tidType(5678)
 	flags := flagsType(syscall.O_RDONLY)
-	cnt1 := Counter{Count: 1, Duration: 1000, DurationToPrev: 100}
+	cnt1 := Counter{Count: 1, Duration: 1000, DurationToPrev: 100, Bytes: 64}
 
 	iod.add(path, traceId, comm, pid, tid, flags, cnt1)
 
 	if iod.paths[path][traceId][comm][pid][tid][flags] != cnt1 {
 		t.Errorf("Expected counter %v, got %v", cnt1, iod.paths[path][traceId][comm][pid][tid][flags])
 	}
-	cnt2 := Counter{Count: 2, Duration: 2000, DurationToPrev: 200}
+	cnt2 := Counter{Count: 2, Duration: 2000, DurationToPrev: 200, Bytes: 128}
 
 	iod.add(path, traceId, comm, pid, tid, flags, cnt2)
 
@@ -41,24 +41,28 @@ func TestMerge(t *testing.T) {
 			Count:          10,
 			Duration:       1000,
 			DurationToPrev: 100,
+			Bytes:          64,
 		}}}}}}}}
 	iod2 := iorData{paths: pathMap{
 		"path1": {traceId: {"comm1": {100: {1000: {roFlag: Counter{
 			Count:          20,
 			Duration:       2000,
 			DurationToPrev: 200,
+			Bytes:          128,
 		}}}}}}}}
 	iod3 := iorData{paths: pathMap{
 		"path2": {traceId: {"comm2": {101: {1000: {roFlag: Counter{
 			Count:          20,
 			Duration:       2000,
 			DurationToPrev: 200,
+			Bytes:          128,
 		}}}}}}}}
 	iod4 := iorData{paths: pathMap{
 		"path2": {traceId: {"comm2": {101: {1000: {roFlag: Counter{
 			Count:          40,
 			Duration:       4000,
 			DurationToPrev: 400,
+			Bytes:          256,
 		}}}}}}}}
 
 	t.Log("iod1", iod1)
@@ -78,6 +82,9 @@ func TestMerge(t *testing.T) {
 		if merged.paths["path2"][traceId]["comm2"][101][1000][roFlag].Count != 60 {
 			t.Errorf("Expected counter 60, got %d", merged.paths["path2"][1]["comm2"][101][1000][roFlag].Count)
 		}
+		if merged.paths["path2"][traceId]["comm2"][101][1000][roFlag].Bytes != 384 {
+			t.Errorf("Expected bytes 384, got %d", merged.paths["path2"][1]["comm2"][101][1000][roFlag].Bytes)
+		}
 	})
 
 	// t.Run("Iterate over lines", func(t *testing.T) {