From 2f95cae8677b5724d6980da4cc77fcc04554bba9 Mon Sep 17 00:00:00 2001
From: Nick Peters <strategicpause@gmail.com>
Date: Sat, 18 Nov 2023 22:00:19 -0800
Subject: [PATCH] Adding flag to instead write random data.

---
 command/memory/leak.go     | 16 ++++++++++++----
 command/memory/register.go | 13 +++++++++----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/command/memory/leak.go b/command/memory/leak.go
index ff3d2bc..b2fb5a4 100644
--- a/command/memory/leak.go
+++ b/command/memory/leak.go
@@ -1,6 +1,7 @@
 package memory
 
 import (
+	"crypto/rand"
 	"fmt"
 	"github.com/strategicpause/memory-leak/metrics"
 	"time"
@@ -11,6 +12,7 @@ type Params struct {
 	BlockSizeInBytes   uint64
 	StepTimeInSeconds  time.Duration
 	PauseTimeInSeconds time.Duration
+	RandomData         bool
 }
 
 func memoryLeak(params *Params) error {
@@ -21,8 +23,13 @@ func memoryLeak(params *Params) error {
 
 	for i := 0; i < numEntries; i++ {
 		list[i] = make([]byte, params.BlockSizeInBytes)
-		for j := 0; j < int(params.BlockSizeInBytes); j++ {
-			list[i][j] = 0
+		if params.RandomData {
+			_, _ = rand.Read(list[i])
+		} else {
+			// This will result in allocating memory from the virtual address space.
+			for j := 0; j < int(params.BlockSizeInBytes); j++ {
+				list[i][j] = 0
+			}
 		}
 		metrics.PrintMemory()
 		time.Sleep(params.StepTimeInSeconds)
@@ -36,6 +43,7 @@ func memoryLeak(params *Params) error {
 }
 
 func PrintParams(params *Params) {
-	fmt.Printf("MaxMemory = %v MiB\tBlockSize = %v MiB\tPauseTime = %v.\n",
-		metrics.BToMiB(params.MaxMemoryInBytes), metrics.BToMiB(params.BlockSizeInBytes), params.PauseTimeInSeconds)
+	fmt.Printf("MaxMemory = %v MiB\tBlockSize = %v MiB\tPauseTime = %v\tRandomData = %t.\n",
+		metrics.BToMiB(params.MaxMemoryInBytes), metrics.BToMiB(params.BlockSizeInBytes),
+		params.PauseTimeInSeconds, params.RandomData)
 }
diff --git a/command/memory/register.go b/command/memory/register.go
index c7fd19a..9af960f 100644
--- a/command/memory/register.go
+++ b/command/memory/register.go
@@ -7,10 +7,11 @@ import (
 )
 
 const (
-	MaxMemoryName     = "max-memory"
-	BlockSizeName     = "block-size"
-	StepDurationName  = "step-duration"
-	PauseDurationName = "pause-duration"
+	MaxMemoryName      = "max-memory"
+	BlockSizeName      = "block-size"
+	StepDurationName   = "step-duration"
+	PauseDurationName  = "pause-duration"
+	WithRandomDataName = "with-random-data"
 )
 
 func Register() cli.Command {
@@ -44,6 +45,9 @@ func flags() []cli.Flag {
 			Usage: "Time to wait, in seconds, after all memory has been allocated.",
 			Value: 0,
 		},
+		cli.BoolFlag{
+			Name: WithRandomDataName,
+		},
 	}
 }
 
@@ -67,6 +71,7 @@ func action(ctx *cli.Context) error {
 		BlockSizeInBytes:   uint64(blockSize),
 		StepTimeInSeconds:  ctx.Duration(StepDurationName),
 		PauseTimeInSeconds: ctx.Duration(PauseDurationName),
+		RandomData:         ctx.Bool(WithRandomDataName),
 	}
 
 	return memoryLeak(params)