Add memory usage benchmarks

Signed-off-by: Sergey Fedorov <sergey.fedorov@protocol.ai>

mir_test.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/filecoin-project/mir"
 	"github.com/filecoin-project/mir/pkg/deploytest"
+	"github.com/filecoin-project/mir/pkg/logging"
 )
 
 var (
@@ -52,6 +53,9 @@ var _ = Describe("Basic test", func() {
 		// When the deployment stops, the final node statuses will be written here.
 		finalStatuses []deploytest.NodeStatus
 
+		heapObjects int64
+		heapAlloc   int64
+
 		ctx context.Context
 	)
 
@@ -90,7 +94,7 @@ var _ = Describe("Basic test", func() {
 		}
 
 		// Run deployment until it stops and returns final node statuses.
-		finalStatuses = deployment.Run(ctx, tickInterval)
+		finalStatuses, heapObjects, heapAlloc = deployment.Run(ctx, tickInterval)
 		fmt.Printf("Deployment run returned.")
 
 		// Check whether all the test replicas exited correctly.
@@ -226,6 +230,30 @@ var _ = Describe("Basic test", func() {
 		}),
 	)
 
+	// This Pending table will be skipped.
+	// Change to FDescribeTable to make it focused.
+	XDescribeTable("Memory usage benchmarks", Serial,
+		func(c *deploytest.TestConfig) {
+			testFunc(c)
+			AddReportEntry("Heap objects", heapObjects)
+			AddReportEntry("Heap allocated (KB)",
+				fmt.Sprintf("%.3f", float64(heapAlloc)/1024))
+		},
+		Entry("Runs for 10s with 4 nodes", &deploytest.TestConfig{
+			NumReplicas: 4,
+			NumClients:  1,
+			Transport:   "fake",
+			Duration:    10 * time.Second,
+			Logger:      logging.ConsoleErrorLogger,
+		}),
+		Entry("Runs for 100s with 4 nodes", &deploytest.TestConfig{
+			NumReplicas: 4,
+			NumClients:  1,
+			Transport:   "fake",
+			Duration:    100 * time.Second,
+			Logger:      logging.ConsoleErrorLogger,
+		}),
+	)
 })
 
 // Remove all temporary data produced by the tests at the end.

pkg/deploytest/deployment.go

@@ -11,6 +11,7 @@ import (
 	"crypto"
 	"fmt"
 	"path/filepath"
+	"runtime"
 	"sync"
 	"time"
 
@@ -194,13 +195,27 @@ func NewDeployment(testConfig *TestConfig) (*Deployment, error) {
 // Run launches the test deployment.
 // It starts all test replicas, the dummy client, and the fake message transport subsystem,
 // waits until the replicas stop, and returns the final statuses of all the replicas.
-func (d *Deployment) Run(ctx context.Context, tickInterval time.Duration) []NodeStatus {
+func (d *Deployment) Run(ctx context.Context, tickInterval time.Duration) (finalStatuses []NodeStatus, heapObjects int64, heapAlloc int64) {
 
 	// Initialize helper variables.
-	finalStatuses := make([]NodeStatus, len(d.TestReplicas))
+	finalStatuses = make([]NodeStatus, len(d.TestReplicas))
 	var nodeWg sync.WaitGroup
 	var clientWg sync.WaitGroup
 
+	ctx2, cancel := context.WithCancel(context.Background())
+
+	var m1, m2 runtime.MemStats
+	runtime.GC()
+	runtime.ReadMemStats(&m1)
+	go func() {
+		<-ctx.Done()
+		runtime.GC()
+		runtime.ReadMemStats(&m2)
+		heapObjects = int64(m2.HeapObjects - m1.HeapObjects)
+		heapAlloc = int64(m2.HeapAlloc - m1.HeapAlloc)
+		cancel()
+	}()
+
 	// Start the Mir nodes.
 	nodeWg.Add(len(d.TestReplicas))
 	for i, testReplica := range d.TestReplicas {
@@ -211,7 +226,7 @@ func (d *Deployment) Run(ctx context.Context, tickInterval time.Duration) []Node
 			defer nodeWg.Done()
 
 			testReplica.Config.Logger.Log(logging.LevelDebug, "running")
-			finalStatuses[i] = testReplica.Run(ctx, tickInterval)
+			finalStatuses[i] = testReplica.Run(ctx2, tickInterval)
 			if err := finalStatuses[i].ExitErr; err != nil {
 				testReplica.Config.Logger.Log(logging.LevelError, "exit with error:", err)
 			} else {
@@ -232,20 +247,20 @@ func (d *Deployment) Run(ctx context.Context, tickInterval time.Duration) []Node
 			defer GinkgoRecover()
 			defer clientWg.Done()
 
-			c.Connect(ctx, d.localRequestReceiverAddrs())
-			submitDummyRequests(ctx, c, d.testConfig.NumNetRequests)
+			c.Connect(ctx2, d.localRequestReceiverAddrs())
+			submitDummyRequests(ctx2, c, d.testConfig.NumNetRequests)
 			c.Disconnect()
 		}(client)
 	}
 
-	<-ctx.Done()
+	<-ctx2.Done()
 
 	// Wait for all replicas and clients to terminate
 	nodeWg.Wait()
 	clientWg.Wait()
 
 	fmt.Printf("All go routines shut down\n")
-	return finalStatuses
+	return
 }
 
 // localGrpcTransport creates an instance of GrpcTransport based on the numeric IDs of test replicas.