From c558daf3e670b713cfaf99d1abb0356c177bb160 Mon Sep 17 00:00:00 2001
From: Nick Peters <strategicpause@gmail.com>
Date: Sun, 17 Sep 2023 20:13:08 -0700
Subject: [PATCH] Cleaning up previous commit.

---
 README.md                  |  5 +++--
 command/socket/leak.go     | 14 ++------------
 command/socket/register.go | 23 ++++++++++++++++++++---
 3 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 81679c2..9ed3dd2 100644
--- a/README.md
+++ b/README.md
@@ -30,9 +30,10 @@ This will turn off swap to make sure OOMs can be recreated at the expected memor
 TCP Leak
 ~~~
 $ podman run \
---memory=1024m \
+--memory=128m \
 --name=tcp-leak \
 --rm \
 localhost/memory-leak \
-./leak tcp
+./leak socket --size 16
 ~~~
+The above command will allocate sockets, establish connections, and write 16 KiB. 
\ No newline at end of file
diff --git a/command/socket/leak.go b/command/socket/leak.go
index 8c496e7..5a9368a 100644
--- a/command/socket/leak.go
+++ b/command/socket/leak.go
@@ -12,22 +12,12 @@ const (
 	// defined by  /proc/sys/net/core/somaxconn.
 	MaxConnections = 4096
 	StartPort      = 9090
-	// By default each established socket connection will write 1 KiB to the buffer.
-	KiB = 1024
 )
 
 var (
 	LocalAddr = [4]byte{127, 0, 0, 1}
 )
 
-type Params struct {
-	NumSockets            int64
-	NetworkAddressDomain  int
-	ConnectionType        int
-	CommunicationProtocol int
-	PauseTimeInSeconds    time.Duration
-}
-
 func tcpLeak(params *Params) error {
 	PrintParams(params)
 
@@ -60,7 +50,7 @@ func tcpLeak(params *Params) error {
 			return err
 		}
 
-		_, err = unix.Write(clientFd, make([]byte, KiB))
+		_, err = unix.Write(clientFd, make([]byte, params.DataSize))
 
 		if i%100 == 0 {
 			metrics.PrintSocketStats()
@@ -106,7 +96,7 @@ func resetServer(params *Params, port int) (chan bool, error) {
 }
 
 func PrintParams(params *Params) {
-	fmt.Printf("NumSockets: %v\n", params.NumSockets)
+	fmt.Printf("NumSockets: %v\tSize: %v\n", params.NumSockets, params.DataSize)
 }
 
 func Must[T any](obj T, err error) T {
diff --git a/command/socket/register.go b/command/socket/register.go
index 5085e21..9454194 100644
--- a/command/socket/register.go
+++ b/command/socket/register.go
@@ -7,9 +7,11 @@ import (
 )
 
 const (
-	NumSocketsName            = "num-sockets"
-	CommunicationProtocolName = "comm-protocol"
-	PauseDurationName         = "pause"
+	NumSocketsName    = "num-sockets"
+	SizeName          = "size"
+	PauseDurationName = "pause"
+
+	KiB = 1024
 )
 
 func Register() cli.Command {
@@ -33,15 +35,30 @@ func flags() []cli.Flag {
 			Usage: "Time between allocations in seconds.",
 			Value: time.Second,
 		},
+		cli.IntFlag{
+			Name:  SizeName,
+			Usage: "Number of KiB to write per per socket.",
+			Value: 4,
+		},
 	}
 }
 
+type Params struct {
+	NumSockets            int64
+	NetworkAddressDomain  int
+	ConnectionType        int
+	CommunicationProtocol int
+	DataSize              int
+	PauseTimeInSeconds    time.Duration
+}
+
 func action(ctx *cli.Context) error {
 	params := &Params{
 		NumSockets:            ctx.Int64(NumSocketsName),
 		NetworkAddressDomain:  unix.AF_INET,
 		ConnectionType:        unix.SOCK_STREAM,
 		CommunicationProtocol: unix.IPPROTO_TCP,
+		DataSize:              ctx.Int(SizeName) * KiB,
 		PauseTimeInSeconds:    ctx.Duration(PauseDurationName),
 	}
 	return tcpLeak(params)