From c3326c440aa7f65355217ef235453587b222de7c Mon Sep 17 00:00:00 2001
From: Zach Koopmans <zkoopmans@google.com>
Date: Fri, 11 Jun 2021 11:52:55 -0700
Subject: [PATCH] [benchmarks] Fix redis benchmark.

Intermittenly, the connection between the client redis_benchmark container can be
flaky with the server, even if the server has been up for a long time. If this
happens, just re-run a client until we get a result.

Also, don't start a new server for each operation.

Also also, modify Makefile run-benchmark call to accept RUNTIME=runc correctly.

PiperOrigin-RevId: 378918886
---
 Makefile                               |  3 +-
 test/benchmarks/database/redis_test.go | 86 ++++++++++++++------------
 2 files changed, 50 insertions(+), 39 deletions(-)

diff --git a/Makefile b/Makefile
index 5aefef0c4..68fb4c892 100644
--- a/Makefile
+++ b/Makefile
@@ -369,7 +369,8 @@ benchmark-platforms: load-benchmarks $(RUNTIME_BIN) ## Runs benchmarks for runc
 .PHONY: benchmark-platforms
 
 run-benchmark: load-benchmarks $(RUNTIME_BIN) ## Runs single benchmark and optionally sends data to BigQuery.
-	@$(call run_benchmark,$(RUNTIME)$(BENCH_VFS),$(BENCH_RUNTIME_ARGS) $(BENCH_VFS))
+	@if test "$(RUNTIME)" = "runc"; then $(call run_benchmark,$(RUNTIME)); fi;
+	@if test "$(RUNTIME)" != "runc"; then $(call run_benchmark,$(RUNTIME)$(BENCH_VFS),$(BENCH_RUNTIME_ARGS) $(BENCH_VFS)); fi;
 .PHONY: run-benchmark
 
 ##
diff --git a/test/benchmarks/database/redis_test.go b/test/benchmarks/database/redis_test.go
index f3c4522ac..7497f69b3 100644
--- a/test/benchmarks/database/redis_test.go
+++ b/test/benchmarks/database/redis_test.go
@@ -1,6 +1,6 @@
-// Copyright 2020 The gVisor Authors.
+// Copyright 2021 The gVisor Authors.
 //
-// Licensed under the Apache License, Version 2.0 (the "License");
+// Licensed under the Apache License, Version 3.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
@@ -17,6 +17,7 @@ package database
 import (
 	"context"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -65,6 +66,34 @@ func BenchmarkRedis(b *testing.B) {
 	// Redis runs on port 6379 by default.
 	port := 6379
 	ctx := context.Background()
+	server := serverMachine.GetContainer(ctx, b)
+	defer server.CleanUp(ctx)
+
+	// The redis docker container takes no arguments to run a redis server.
+	if err := server.Spawn(ctx, dockerutil.RunOpts{
+		Image: "benchmarks/redis",
+		Ports: []int{port},
+	}); err != nil {
+		b.Fatalf("failed to start redis server with: %v", err)
+	}
+
+	if out, err := server.WaitForOutput(ctx, "Ready to accept connections", 3*time.Second); err != nil {
+		b.Fatalf("failed to start redis server: %v %s", err, out)
+	}
+
+	ip, err := serverMachine.IPAddress()
+	if err != nil {
+		b.Fatalf("failed to get IP from server: %v", err)
+	}
+
+	serverPort, err := server.FindPort(ctx, port)
+	if err != nil {
+		b.Fatalf("failed to get IP from server: %v", err)
+	}
+
+	if err = harness.WaitUntilServing(ctx, clientMachine, ip, serverPort); err != nil {
+		b.Fatalf("failed to start redis with: %v", err)
+	}
 	for _, operation := range operations {
 		param := tools.Parameter{
 			Name:  "operation",
@@ -74,49 +103,30 @@ func BenchmarkRedis(b *testing.B) {
 		if err != nil {
 			b.Fatalf("Failed to parse paramaters: %v", err)
 		}
+
 		b.Run(name, func(b *testing.B) {
-			server := serverMachine.GetContainer(ctx, b)
-			defer server.CleanUp(ctx)
-
-			// The redis docker container takes no arguments to run a redis server.
-			if err := server.Spawn(ctx, dockerutil.RunOpts{
-				Image: "benchmarks/redis",
-				Ports: []int{port},
-			}); err != nil {
-				b.Fatalf("failed to start redis server with: %v", err)
-			}
-
-			if out, err := server.WaitForOutput(ctx, "Ready to accept connections", 3*time.Second); err != nil {
-				b.Fatalf("failed to start redis server: %v %s", err, out)
-			}
-
-			ip, err := serverMachine.IPAddress()
-			if err != nil {
-				b.Fatalf("failed to get IP from server: %v", err)
-			}
-
-			serverPort, err := server.FindPort(ctx, port)
-			if err != nil {
-				b.Fatalf("failed to get IP from server: %v", err)
-			}
-
-			if err = harness.WaitUntilServing(ctx, clientMachine, ip, serverPort); err != nil {
-				b.Fatalf("failed to start redis with: %v", err)
-			}
-
-			client := clientMachine.GetNativeContainer(ctx, b)
-			defer client.CleanUp(ctx)
-
 			redis := tools.Redis{
 				Operation: operation,
 			}
-			b.ResetTimer()
-			out, err := client.Run(ctx, dockerutil.RunOpts{
-				Image: "benchmarks/redis",
-			}, redis.MakeCmd(ip, serverPort, b.N /*requests*/)...)
+
+			// Sometimes, the connection between the redis client and server can be
+			// flaky such that the client returns infinity as the QPS measurement for
+			// a give operation. If this happens, retry the client up to 3 times.
+			out := "inf"
+			for retries := 0; strings.Contains(out, "inf") && retries < 3; retries++ {
+				b.ResetTimer()
+				client := clientMachine.GetNativeContainer(ctx, b)
+				defer client.CleanUp(ctx)
+
+				out, err = client.Run(ctx, dockerutil.RunOpts{
+					Image: "benchmarks/redis",
+				}, redis.MakeCmd(ip, serverPort, b.N /*requests*/)...)
+			}
+
 			if err != nil {
 				b.Fatalf("redis-benchmark failed with: %v", err)
 			}
+
 			b.StopTimer()
 			redis.Report(b, out)
 		})