add support for --store=type=mem,size=val when start a server

This commit sets `--store=type=mem` (storing database in memory)
as the default option for storage.
We also allow users to customize the proportion of available
memory allocated to the test server.

Author: ZhouXing19

testserver/testserver.go

@@ -56,10 +56,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cockroachdb/cockroach-go/v2/testserver/version"
 	// Import postgres driver.
 	_ "github.com/lib/pq"
-
-	"github.com/cockroachdb/cockroach-go/v2/testserver/version"
 )
 
 var customBinaryFlag = flag.String("cockroach-binary", "", "Use specified cockroach binary")
@@ -77,6 +76,9 @@ const (
 	firstTenantID = 2
 )
 
+// By default, we allocate 20% of available memory to the test server.
+const defaultStoreMemSize = 0.2
+
 // TestServer is a helper to run a real cockroach node.
 type TestServer interface {
 	// Start starts the server.
@@ -166,8 +168,10 @@ func NewDBForTestWithDatabase(
 type TestServerOpt func(args *testServerArgs)
 
 type testServerArgs struct {
-	secure bool
-	rootPW string // if nonempty, set as pw for root
+	secure       bool
+	rootPW       string  // if nonempty, set as pw for root
+	storeOnDisk  bool    // to save database in disk
+	storeMemSize float64 // the proportion of available memory allocated to test server
 }
 
 // SecureOpt is a TestServer option that can be passed to NewTestServer to
@@ -178,6 +182,27 @@ func SecureOpt() TestServerOpt {
 	}
 }
 
+// StoreOnDiskOpt is a TestServer option that can be passed to NewTestServer
+// to enable storing database in memory.
+func StoreOnDiskOpt() TestServerOpt {
+	return func(args *testServerArgs) {
+		args.storeOnDisk = true
+	}
+}
+
+// SetStoreMemSizeOpt is a TestServer option that can be passed to NewTestServer
+// to set the proportion of available memory that is allocated
+// to the test server.
+func SetStoreMemSizeOpt(memSize float64) TestServerOpt {
+	return func(args *testServerArgs) {
+		if memSize > 0 {
+			args.storeMemSize = memSize
+		} else {
+			args.storeMemSize = defaultStoreMemSize
+		}
+	}
+}
+
 // RootPasswordOpt is a TestServer option that, when passed to NewTestServer,
 // sets the given password for the root user (and returns a URL using it from
 // PGURL(). This avoids having to use client certs.
@@ -200,6 +225,7 @@ const (
 // found in your path.
 func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 	serverArgs := &testServerArgs{}
+	serverArgs.storeMemSize = defaultStoreMemSize
 	for _, applyOptToArgs := range opts {
 		applyOptToArgs(serverArgs)
 	}
@@ -236,6 +262,9 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 		}
 		return path, nil
 	}
+	// TODO(janexing): Make sure the log is written to logDir instead of shown in console.
+	// Should be done once issue #109 is solved:
+	// https://github.com/cockroachdb/cockroach-go/issues/109
 	logDir, err := mkDir(logsDirName)
 	if err != nil {
 		return nil, err
@@ -292,6 +321,13 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 		startCmd = "start"
 	}
 
+	var storeArg string
+	if serverArgs.storeOnDisk {
+		storeArg = "--store=path=" + baseDir
+	} else {
+		storeArg = fmt.Sprintf("--store=type=mem,size=%.2f", serverArgs.storeMemSize)
+	}
+
 	args := []string{
 		cockroachBinary,
 		startCmd,
@@ -300,7 +336,7 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 		"--host=localhost",
 		"--port=0",
 		"--http-port=0",
-		"--store=" + baseDir,
+		storeArg,
 		"--listening-url-file=" + listeningURLFile,
 	}
 

testserver/testserver_test.go

@@ -24,6 +24,8 @@ import (
 
 const noPW = ""
 
+const defStoreMemSize = 0.2
+
 func TestRunServer(t *testing.T) {
 	const testPW = "foobar"
 	for _, tc := range []struct {
@@ -34,6 +36,12 @@ func TestRunServer(t *testing.T) {
 			name:          "Insecure",
 			instantiation: func(t *testing.T) (*sql.DB, func()) { return testserver.NewDBForTest(t) },
 		},
+		{
+			name: "InsecureWithCustomizedMemSize",
+			instantiation: func(t *testing.T) (*sql.DB, func()) {
+				return testserver.NewDBForTest(t, testserver.SetStoreMemSizeOpt(0.3))
+			},
+		},
 		{
 			name:          "SecureClientCert",
 			instantiation: func(t *testing.T) (*sql.DB, func()) { return testserver.NewDBForTest(t, testserver.SecureOpt()) },
@@ -44,34 +52,46 @@ func TestRunServer(t *testing.T) {
 				return testserver.NewDBForTest(t, testserver.SecureOpt(), testserver.RootPasswordOpt(testPW))
 			},
 		},
+		{
+			name: "InsecureTenantStoreOnDisk",
+			instantiation: func(t *testing.T) (*sql.DB, func()) {
+				return testserver.NewDBForTest(t, testserver.StoreOnDiskOpt())
+			},
+		},
+		{
+			name: "SecureTenantStoreOnDisk",
+			instantiation: func(t *testing.T) (*sql.DB, func()) {
+				return testserver.NewDBForTest(t, testserver.SecureOpt(), testserver.StoreOnDiskOpt())
+			},
+		},
 		{
 			name: "InsecureTenant",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
-				return newTenantDBForTest(t, false /* secure */, false /* proxy */, noPW)
+				return newTenantDBForTest(t, false /* secure */, false /* proxy */, noPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 			},
 		},
 		{
 			name: "SecureTenant",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
-				return newTenantDBForTest(t, true /* secure */, false /* proxy */, noPW)
+				return newTenantDBForTest(t, true /* secure */, false /* proxy */, noPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 			},
 		},
 		{
 			name: "SecureTenantCustomPassword",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
-				return newTenantDBForTest(t, true /* secure */, false /* proxy */, testPW)
+				return newTenantDBForTest(t, true /* secure */, false /* proxy */, testPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 			},
 		},
 		{
 			name: "SecureTenantThroughProxy",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
-				return newTenantDBForTest(t, true /* secure */, true /* proxy */, noPW)
+				return newTenantDBForTest(t, true /* secure */, true /* proxy */, noPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 			},
 		},
 		{
 			name: "SecureTenantThroughProxyCustomPassword",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
-				return newTenantDBForTest(t, true /* secure */, true /* proxy */, testPW)
+				return newTenantDBForTest(t, true /* secure */, true /* proxy */, testPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 			},
 		},
 	} {
@@ -105,15 +125,25 @@ type tenantInterface interface {
 // newTenantDBForTest is a testing helper function that starts a TestServer
 // process and a SQL tenant process pointed at this TestServer. A sql connection
 // to the tenant and a cleanup function are returned.
-func newTenantDBForTest(t *testing.T, secure bool, proxy bool, pw string) (*sql.DB, func()) {
+func newTenantDBForTest(
+	t *testing.T, secure bool, proxy bool, pw string, diskStore bool, storeMemSize float64,
+) (*sql.DB, func()) {
 	t.Helper()
 	var opts []testserver.TestServerOpt
 	if secure {
 		opts = append(opts, testserver.SecureOpt())
 	}
+	if diskStore {
+		opts = append(opts, testserver.StoreOnDiskOpt())
+	}
 	if pw != "" {
 		opts = append(opts, testserver.RootPasswordOpt(pw))
 	}
+	if storeMemSize >= 0 {
+		opts = append(opts, testserver.SetStoreMemSizeOpt(storeMemSize))
+	} else {
+		t.Fatal("Percentage memory size for data storage cannot be nagative")
+	}
 	ts, err := testserver.NewTestServer(opts...)
 	if err != nil {
 		t.Fatal(err)
@@ -138,7 +168,7 @@ func newTenantDBForTest(t *testing.T, secure bool, proxy bool, pw string) (*sql.
 }
 
 func TestTenant(t *testing.T) {
-	db, stop := newTenantDBForTest(t, false /* secure */, false /* proxy */, noPW)
+	db, stop := newTenantDBForTest(t, false /* secure */, false /* proxy */, noPW, false /* diskStore */, defStoreMemSize /* storeMemSize */)
 	defer stop()
 	if _, err := db.Exec("SELECT 1"); err != nil {
 		t.Fatal(err)