From fe6e3863ba322e0ef19f0e75bf675ab5f2a0bb6d Mon Sep 17 00:00:00 2001
From: Mike Poindexter <mpoindexter@housecanary.com>
Date: Wed, 8 Apr 2020 11:36:37 -0700
Subject: [PATCH] Add bench tests for knn

---
 tests/keys_search_test.go | 17 +++++++
 tests/mock_test.go        | 16 +++++--
 tests/tests_test.go       | 97 +++++++++++++++++++++++++++++++++++++--
 3 files changed, 121 insertions(+), 9 deletions(-)

diff --git a/tests/keys_search_test.go b/tests/keys_search_test.go
index ff236d81..e78ae726 100644
--- a/tests/keys_search_test.go
+++ b/tests/keys_search_test.go
@@ -2,6 +2,7 @@ package tests
 
 import (
 	"fmt"
+	"math/rand"
 	"sort"
 	"testing"
 )
@@ -429,3 +430,19 @@ func match(expectIn string) func(org, v interface{}) (resp, expect interface{})
 		return fmt.Sprintf("%v", org), expectIn
 	}
 }
+
+func subBenchSearch(b *testing.B, mc *mockServer) {
+	runBenchStep(b, mc, "KNN", keys_KNN_bench)
+}
+
+func keys_KNN_bench(mc *mockServer) error {
+	lat := rand.Float64()*180 - 90
+	lon := rand.Float64()*360 - 180
+	_, err := mc.conn.Do("NEARBY",
+		"mykey",
+		"LIMIT", 50,
+		"DISTANCE",
+		"POINTS",
+		"POINT", lat, lon)
+	return err
+}
diff --git a/tests/mock_test.go b/tests/mock_test.go
index 77825e87..fb8a1037 100644
--- a/tests/mock_test.go
+++ b/tests/mock_test.go
@@ -20,15 +20,19 @@ import (
 
 var errTimeout = errors.New("timeout")
 
-func mockCleanup() {
-	fmt.Printf("Cleanup: may take some time... ")
+func mockCleanup(silent bool) {
+	if !silent {
+		fmt.Printf("Cleanup: may take some time... ")
+	}
 	files, _ := ioutil.ReadDir(".")
 	for _, file := range files {
 		if strings.HasPrefix(file.Name(), "data-mock-") {
 			os.RemoveAll(file.Name())
 		}
 	}
-	fmt.Printf("OK\n")
+	if !silent {
+		fmt.Printf("OK\n")
+	}
 }
 
 type mockServer struct {
@@ -39,11 +43,13 @@ type mockServer struct {
 	conn redis.Conn
 }
 
-func mockOpenServer() (*mockServer, error) {
+func mockOpenServer(silent bool) (*mockServer, error) {
 	rand.Seed(time.Now().UnixNano())
 	port := rand.Int()%20000 + 20000
 	dir := fmt.Sprintf("data-mock-%d", port)
-	fmt.Printf("Starting test server at port %d\n", port)
+	if !silent {
+		fmt.Printf("Starting test server at port %d\n", port)
+	}
 	logOutput := ioutil.Discard
 	if os.Getenv("PRINTLOG") == "1" {
 		logOutput = os.Stderr
diff --git a/tests/tests_test.go b/tests/tests_test.go
index 34953ddf..44d347dd 100644
--- a/tests/tests_test.go
+++ b/tests/tests_test.go
@@ -2,10 +2,14 @@ package tests
 
 import (
 	"fmt"
+	"math/rand"
 	"os"
 	"os/signal"
 	"syscall"
 	"testing"
+	"time"
+
+	"github.com/gomodule/redigo/redis"
 )
 
 const (
@@ -23,18 +27,18 @@ const (
 )
 
 func TestAll(t *testing.T) {
-	mockCleanup()
-	defer mockCleanup()
+	mockCleanup(false)
+	defer mockCleanup(false)
 
 	ch := make(chan os.Signal)
 	signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		<-ch
-		mockCleanup()
+		mockCleanup(false)
 		os.Exit(1)
 	}()
 
-	mc, err := mockOpenServer()
+	mc, err := mockOpenServer(false)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -83,3 +87,88 @@ func runStep(t *testing.T, mc *mockServer, name string, step func(mc *mockServer
 		fmt.Printf("["+green+"ok"+clear+"]: %s\n", name)
 	})
 }
+
+func BenchmarkAll(b *testing.B) {
+	mockCleanup(true)
+	defer mockCleanup(true)
+
+	ch := make(chan os.Signal)
+	signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-ch
+		mockCleanup(true)
+		os.Exit(1)
+	}()
+
+	mc, err := mockOpenServer(true)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer mc.Close()
+	runSubBenchmark(b, "search", mc, subBenchSearch)
+}
+
+func loadBenchmarkPoints(b *testing.B, mc *mockServer) (err error) {
+	const nPoints = 200000
+	rand.Seed(time.Now().UnixNano())
+
+	// add a bunch of points
+	for i := 0; i < nPoints; i++ {
+		val := fmt.Sprintf("val:%d", i)
+		var resp string
+		var lat, lon, fval float64
+		fval = rand.Float64()
+		lat = rand.Float64()*180 - 90
+		lon = rand.Float64()*360 - 180
+		resp, err = redis.String(mc.conn.Do("SET",
+			"mykey", val,
+			"FIELD", "foo", fval,
+			"POINT", lat, lon))
+		if err != nil {
+			return
+		}
+		if resp != "OK" {
+			err = fmt.Errorf("expected 'OK', got '%s'", resp)
+			return
+		}
+	}
+	return
+}
+
+func runSubBenchmark(b *testing.B, name string, mc *mockServer, bench func(t *testing.B, mc *mockServer)) {
+	b.Run(name, func(b *testing.B) {
+		bench(b, mc)
+	})
+}
+
+func runBenchStep(b *testing.B, mc *mockServer, name string, step func(mc *mockServer) error) {
+	b.Helper()
+	b.Run(name, func(b *testing.B) {
+		b.Helper()
+		if err := func() error {
+			// reset the current server
+			mc.ResetConn()
+			defer mc.ResetConn()
+			// clear the database so the test is consistent
+			if err := mc.DoBatch([][]interface{}{
+				{"OUTPUT", "resp"}, {"OK"},
+				{"FLUSHDB"}, {"OK"},
+			}); err != nil {
+				return err
+			}
+			err := loadBenchmarkPoints(b, mc)
+			if err != nil {
+				return err
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				if err := step(mc); err != nil {
+					return err
+				}
+			}
+			return nil
+		}(); err != nil {
+			b.Fatal(err)
+		}
+	})
+}