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>

7 files changed, 192 insertions, 42 deletions

diff --git a/internal/event/pair.go b/internal/event/pair.go
index d67ea06..1e9a544 100644
--- a/internal/event/pair.go
+++ b/internal/event/pair.go
@@ -15,6 +15,7 @@ type Pair struct {
 	Comm            string
 	Duration        uint64
 	DurationToPrev  uint64
+	Bytes           uint64 // Number of bytes transferred (read/write/transfer syscalls only)
 	Equals          bool
 }
 
@@ -26,6 +27,7 @@ func NewPair(enterEv Event) *Pair {
 	e.Comm = ""
 	e.Duration = 0
 	e.DurationToPrev = 0
+	e.Bytes = 0
 	e.Equals = false
 	return e
 }
@@ -99,6 +101,7 @@ func (e *Pair) Recycle() {
 	e.Comm = ""
 	e.Duration = 0
 	e.DurationToPrev = 0
+	e.Bytes = 0
 	e.Equals = false
 	poolOfEventPairs.Put(e)
 }
diff --git a/internal/eventloop.go b/internal/eventloop.go
index 8d5a8b0..169f20f 100644
--- a/internal/eventloop.go
+++ b/internal/eventloop.go
@@ -20,16 +20,15 @@ import (
 
 const sysEnterNameToHandleAtName = "name_to_handle_at"
 
-// TOOD: read and write syscalls: can also collect amount of bytes!
 type eventLoop struct {
-	filter        *eventFilter
-	enterEvs      map[uint32]*event.Pair      // Temp. store of sys_enter tracepoints per Tid.
+	filter         *eventFilter
+	enterEvs       map[uint32]*event.Pair      // Temp. store of sys_enter tracepoints per Tid.
 	pendingHandles map[uint32]string           // map of TID to pathname from name_to_handle_at
-	files         map[int32]file.File         // Track all open files by file descriptor..
-	comms         map[uint32]string           // Program or thread name of the current Tid.
-	prevPairTimes map[uint32]uint64           // Previous event's time (to calculate time differences between two events)
-	flamegraph    flamegraph.IorDataCollector // Storing all paths in a map structure for analysis
-	printCb       func(ep *event.Pair)        // Callback to print the event
+	files          map[int32]file.File         // Track all open files by file descriptor..
+	comms          map[uint32]string           // Program or thread name of the current Tid.
+	prevPairTimes  map[uint32]uint64           // Previous event's time (to calculate time differences between two events)
+	flamegraph     flamegraph.IorDataCollector // Storing all paths in a map structure for analysis
+	printCb        func(ep *event.Pair)        // Callback to print the event
 
 	// Statistics
 	numTracepoints          uint
@@ -42,15 +41,15 @@ type eventLoop struct {
 
 func newEventLoop() *eventLoop {
 	return &eventLoop{
-		filter:        newEventFilter(),
-		enterEvs:      make(map[uint32]*event.Pair),
+		filter:         newEventFilter(),
+		enterEvs:       make(map[uint32]*event.Pair),
 		pendingHandles: make(map[uint32]string),
-		files:         make(map[int32]file.File),
-		comms:         make(map[uint32]string),
-		prevPairTimes: make(map[uint32]uint64),
-		printCb:       func(ep *event.Pair) { fmt.Println(ep); ep.Recycle() },
-		flamegraph:    flamegraph.New(),
-		done:          make(chan struct{}),
+		files:          make(map[int32]file.File),
+		comms:          make(map[uint32]string),
+		prevPairTimes:  make(map[uint32]uint64),
+		printCb:        func(ep *event.Pair) { fmt.Println(ep); ep.Recycle() },
+		flamegraph:     flamegraph.New(),
+		done:           make(chan struct{}),
 	}
 }
 
@@ -296,6 +295,10 @@ func (e *eventLoop) tracepointExited(exitEv event.Event, ch chan<- *event.Pair)
 			}
 		}
 
+		if retEv, ok := ep.ExitEv.(*RetEvent); ok {
+			ep.Bytes = bytesFromRet(retEv)
+		}
+
 	case *Dup3Event:
 		dup3Event := ep.EnterEv.(*Dup3Event)
 		fd := int32(dup3Event.Fd)
@@ -437,3 +440,17 @@ func (e *eventLoop) comm(tid uint32) string {
 	}
 	return ""
 }
+
+// bytesFromRet extracts the number of bytes transferred from a RetEvent.
+// Returns 0 for nil events, errors (Ret <= 0), or unclassified syscalls.
+func bytesFromRet(retEv *types.RetEvent) uint64 {
+	if retEv == nil || retEv.Ret <= 0 {
+		return 0
+	}
+	switch retEv.RetType {
+	case types.READ_CLASSIFIED, types.WRITE_CLASSIFIED, types.TRANSFER_CLASSIFIED:
+		return uint64(retEv.Ret)
+	default:
+		return 0
+	}
+}
diff --git a/internal/eventloop_bytes_test.go b/internal/eventloop_bytes_test.go
new file mode 100644
index 0000000..ed7f7af
--- /dev/null
+++ b/internal/eventloop_bytes_test.go
@@ -0,0 +1,73 @@
+package internal
+
+import (
+	"testing"
+
+	"ior/internal/types"
+)
+
+func TestBytesFromRet(t *testing.T) {
+	tests := []struct {
+		name     string
+		retEvent *types.RetEvent
+		expected uint64
+	}{
+		{name: "nil", retEvent: nil, expected: 0},
+		{
+			name: "negative",
+			retEvent: &types.RetEvent{
+				Ret:     -1,
+				RetType: types.READ_CLASSIFIED,
+			},
+			expected: 0,
+		},
+		{
+			name: "zero",
+			retEvent: &types.RetEvent{
+				Ret:     0,
+				RetType: types.READ_CLASSIFIED,
+			},
+			expected: 0,
+		},
+		{
+			name: "unclassified",
+			retEvent: &types.RetEvent{
+				Ret:     512,
+				RetType: types.UNCLASSIFIED,
+			},
+			expected: 0,
+		},
+		{
+			name: "read",
+			retEvent: &types.RetEvent{
+				Ret:     128,
+				RetType: types.READ_CLASSIFIED,
+			},
+			expected: 128,
+		},
+		{
+			name: "write",
+			retEvent: &types.RetEvent{
+				Ret:     256,
+				RetType: types.WRITE_CLASSIFIED,
+			},
+			expected: 256,
+		},
+		{
+			name: "transfer",
+			retEvent: &types.RetEvent{
+				Ret:     1024,
+				RetType: types.TRANSFER_CLASSIFIED,
+			},
+			expected: 1024,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := bytesFromRet(tt.retEvent); got != tt.expected {
+				t.Errorf("bytesFromRet() = %d, want %d", got, tt.expected)
+			}
+		})
+	}
+}
diff --git a/internal/eventloop_test.go b/internal/eventloop_test.go
index 3b9cb9a..7daa32b 100644
--- a/internal/eventloop_test.go
+++ b/internal/eventloop_test.go
@@ -30,13 +30,13 @@ func TestEventloop(t *testing.T) {
 		"OpenEventTest2": makeOpenEventTestData2(t),
 		"OpenEventTest3": makeOpenEventTestData3(t),
 		// FdEvent tests
-		"ReadEventTest":      makeReadEventTestData(t),
-		"WriteEventTest":     makeWriteEventTestData(t),
-		"CloseEventTest":     makeCloseEventTestData(t),
-		"CloseRangeEventTest": makeCloseRangeEventTestData(t),
+		"ReadEventTest":         makeReadEventTestData(t),
+		"WriteEventTest":        makeWriteEventTestData(t),
+		"CloseEventTest":        makeCloseEventTestData(t),
+		"CloseRangeEventTest":   makeCloseRangeEventTestData(t),
 		"CloseRangeFailureTest": makeCloseRangeFailureTestData(t),
-		"FsyncEventTest":     makeFsyncEventTestData(t),
-		"FtruncateEventTest": makeFtruncateEventTestData(t),
+		"FsyncEventTest":        makeFsyncEventTestData(t),
+		"FtruncateEventTest":    makeFtruncateEventTestData(t),
 		// PathEvent tests
 		"MkdirEventTest":  makeMkdirEventTestData(t),
 		"UnlinkEventTest": makeUnlinkEventTestData(t),
@@ -48,22 +48,22 @@ func TestEventloop(t *testing.T) {
 		"LinkEventTest":    makeLinkEventTestData(t),
 		"SymlinkEventTest": makeSymlinkEventTestData(t),
 		// NullEvent tests
-		"SyncEventTest":         makeSyncEventTestData(t),
-		"IoUringSetupEventTest": makeIoUringSetupEventTestData(t),
-		"IoUringSetupFailureTest": makeIoUringSetupFailureTestData(t),
-		"IoUringEnterEventTest": makeIoUringEnterEventTestData(t),
+		"SyncEventTest":            makeSyncEventTestData(t),
+		"IoUringSetupEventTest":    makeIoUringSetupEventTestData(t),
+		"IoUringSetupFailureTest":  makeIoUringSetupFailureTestData(t),
+		"IoUringEnterEventTest":    makeIoUringEnterEventTestData(t),
 		"IoUringRegisterEventTest": makeIoUringRegisterEventTestData(t),
 		// Dup3Event tests
 		"Dup3EventTest":       makeDup3EventTestData(t),
 		"Dup3WithCloexecTest": makeDup3WithCloexecTestData(t),
 		"Dup2Test":            makeDup2TestData(t),
 		// FcntlEvent tests
-		"FcntlSetFlagsTest":     makeFcntlSetFlagsTestData(t),
-		"FcntlDupfdTest":        makeFcntlDupfdTestData(t),
-		"FcntlDupfdCloexecTest": makeFcntlDupfdCloexecTestData(t),
-		"FcntlErrorTest":        makeFcntlErrorTestData(t),
-		"FcntlInvalidFdTest":    makeFcntlInvalidFdTestData(t),
-		"NameToHandleAtTest":    makeNameToHandleAtTestData(t),
+		"FcntlSetFlagsTest":         makeFcntlSetFlagsTestData(t),
+		"FcntlDupfdTest":            makeFcntlDupfdTestData(t),
+		"FcntlDupfdCloexecTest":     makeFcntlDupfdCloexecTestData(t),
+		"FcntlErrorTest":            makeFcntlErrorTestData(t),
+		"FcntlInvalidFdTest":        makeFcntlInvalidFdTestData(t),
+		"NameToHandleAtTest":        makeNameToHandleAtTestData(t),
 		"NameToHandleAtFailureTest": makeNameToHandleAtFailureTestData(t),
 		// FD Lifecycle tests
 		"FdLifecycleTest": makeFdLifecycleTestData(t),
@@ -304,15 +304,40 @@ func makeReadEventTestData(t *testing.T) (td testData) {
 	enterEv, enterEvBytes := makeEnterFdEvent(t, defaulTime, defaultPid, defaultTid, fd, types.SYS_ENTER_READ)
 	td.rawTracepoints = append(td.rawTracepoints, enterEvBytes)
 
-	exitEv, exitEvBytes := makeExitFdEvent(t, defaulTime+100, defaultPid, defaultTid, fd, types.SYS_EXIT_READ)
+	retBytes := int64(128)
+	_, exitEvBytes := makeExitFdEvent(t, defaulTime+100, defaultPid, defaultTid, fd, types.SYS_EXIT_READ)
+	retEvent := types.RetEvent{
+		EventType: types.EXIT_RET_EVENT,
+		TraceId:   types.SYS_EXIT_READ,
+		Time:      defaulTime + 100,
+		Ret:       retBytes,
+		Pid:       defaultPid,
+		Tid:       defaultTid,
+		RetType:   types.READ_CLASSIFIED,
+	}
+	retEventBytes, err := retEvent.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+	exitEvBytes = retEventBytes
 	td.rawTracepoints = append(td.rawTracepoints, exitEvBytes)
 
 	td.validates = append(td.validates, func(t *testing.T, el *eventLoop, ep *event.Pair) {
 		if !enterEv.Equals(ep.EnterEv) {
 			t.Errorf("Expected '%v' but got '%v'", enterEv, ep.EnterEv)
 		}
-		if !exitEv.Equals(ep.ExitEv) {
-			t.Errorf("Expected '%v' but got '%v'", exitEv, ep.ExitEv)
+		if !retEvent.Equals(ep.ExitEv) {
+			t.Errorf("Expected '%v' but got '%v'", retEvent, ep.ExitEv)
+		}
+		retEv, ok := ep.ExitEv.(*types.RetEvent)
+		if !ok {
+			t.Fatalf("Expected exit event to be *types.RetEvent")
+		}
+		if retEv.Ret != retBytes {
+			t.Errorf("Expected ret bytes %d but got %d", retBytes, retEv.Ret)
+		}
+		if ep.Bytes != uint64(retBytes) {
+			t.Errorf("Expected bytes %d but got %d", retBytes, ep.Bytes)
 		}
 	})
 
@@ -324,15 +349,40 @@ func makeWriteEventTestData(t *testing.T) (td testData) {
 	enterEv, enterEvBytes := makeEnterFdEvent(t, defaulTime, defaultPid, defaultTid, fd, types.SYS_ENTER_WRITE)
 	td.rawTracepoints = append(td.rawTracepoints, enterEvBytes)
 
-	exitEv, exitEvBytes := makeExitFdEvent(t, defaulTime+100, defaultPid, defaultTid, fd, types.SYS_EXIT_WRITE)
+	retBytes := int64(256)
+	_, exitEvBytes := makeExitFdEvent(t, defaulTime+100, defaultPid, defaultTid, fd, types.SYS_EXIT_WRITE)
+	retEvent := types.RetEvent{
+		EventType: types.EXIT_RET_EVENT,
+		TraceId:   types.SYS_EXIT_WRITE,
+		Time:      defaulTime + 100,
+		Ret:       retBytes,
+		Pid:       defaultPid,
+		Tid:       defaultTid,
+		RetType:   types.WRITE_CLASSIFIED,
+	}
+	retEventBytes, err := retEvent.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+	exitEvBytes = retEventBytes
 	td.rawTracepoints = append(td.rawTracepoints, exitEvBytes)
 
 	td.validates = append(td.validates, func(t *testing.T, el *eventLoop, ep *event.Pair) {
 		if !enterEv.Equals(ep.EnterEv) {
 			t.Errorf("Expected '%v' but got '%v'", enterEv, ep.EnterEv)
 		}
-		if !exitEv.Equals(ep.ExitEv) {
-			t.Errorf("Expected '%v' but got '%v'", exitEv, ep.ExitEv)
+		if !retEvent.Equals(ep.ExitEv) {
+			t.Errorf("Expected '%v' but got '%v'", retEvent, ep.ExitEv)
+		}
+		retEv, ok := ep.ExitEv.(*types.RetEvent)
+		if !ok {
+			t.Fatalf("Expected exit event to be *types.RetEvent")
+		}
+		if retEv.Ret != retBytes {
+			t.Errorf("Expected ret bytes %d but got %d", retBytes, retEv.Ret)
+		}
+		if ep.Bytes != uint64(retBytes) {
+			t.Errorf("Expected bytes %d but got %d", retBytes, ep.Bytes)
 		}
 	})
 
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) {