diff --git a/testserver/tenant.go b/testserver/tenant.go index e22bbf5..709b85f 100644 --- a/testserver/tenant.go +++ b/testserver/tenant.go @@ -93,7 +93,7 @@ func (ts *testServerImpl) NewTenantServer(proxy bool) (TestServer, error) { } secureFlag := "--insecure" - certsDir := filepath.Join(ts.baseDir, "certs") + certsDir := filepath.Join(ts.BaseDir(), "certs") if ts.serverArgs.secure { secureFlag = "--certs-dir=" + certsDir // Create tenant client certificate. @@ -235,8 +235,8 @@ func (ts *testServerImpl) NewTenantServer(proxy bool) (TestServer, error) { // TODO(asubiotto): Specify listeningURLFile once we support dynamic // ports. listeningURLFile: "", - stdout: filepath.Join(ts.baseDir, logsDirName, fmt.Sprintf("cockroach.tenant.%d.stdout", tenantID)), - stderr: filepath.Join(ts.baseDir, logsDirName, fmt.Sprintf("cockroach.tenant.%d.stderr", tenantID)), + stdout: filepath.Join(ts.BaseDir(), logsDirName, fmt.Sprintf("cockroach.tenant.%d.stdout", tenantID)), + stderr: filepath.Join(ts.BaseDir(), logsDirName, fmt.Sprintf("cockroach.tenant.%d.stderr", tenantID)), }, } @@ -244,7 +244,6 @@ func (ts *testServerImpl) NewTenantServer(proxy bool) (TestServer, error) { serverArgs: ts.serverArgs, version: ts.version, serverState: stateNew, - baseDir: ts.baseDir, nodes: nodes, } diff --git a/testserver/testserver.go b/testserver/testserver.go index c476eae..cb2c529 100644 --- a/testserver/testserver.go +++ b/testserver/testserver.go @@ -152,7 +152,6 @@ type testServerImpl struct { version *version.Version serverArgs testServerArgs serverState int - baseDir string pgURL []pgURLChan initCmd *exec.Cmd initCmdArgs []string @@ -240,6 +239,8 @@ type testServerArgs struct { envVars []string // to be passed to cmd.Env localityFlags []string cockroachLogsDir string + noFileCleanup bool // do not clean files at `Stop` + baseDir string } // CockroachBinaryPathOpt is a TestServer option that can be passed to @@ -274,6 +275,23 @@ func StoreOnDiskOpt() TestServerOpt { } } +// NoFileCleanup is a TestServer option that can be passed to NewTestServer +// to skip cleanup of files when Testserver is stopped +func NoFileCleanup() TestServerOpt { + return func(args *testServerArgs) { + args.noFileCleanup = true + } +} + +// BasedirOverride is a TestServer option that can be passed to NewTestServer +// to use a given path as testserver base-dir (as opposed to creating a +// temporary one on the fly, which is the default behavior). +func BasedirOverride(baseDirPath string) TestServerOpt { + return func(args *testServerArgs) { + args.baseDir = baseDirPath + } +} + // 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. @@ -428,20 +446,26 @@ var errStoppedInMiddle = errors.New("download stopped in middle") // If the download fails, we attempt just call "cockroach", hoping it is // found in your path. func NewTestServer(opts ...TestServerOpt) (TestServer, error) { - baseDir, err := os.MkdirTemp("", "cockroach-testserver") - if err != nil { - return nil, fmt.Errorf("%s: could not create temp directory: %w", testserverMessagePrefix, err) - } - serverArgs := &testServerArgs{numNodes: 1} serverArgs.storeMemSize = defaultStoreMemSize serverArgs.initTimeoutSeconds = defaultInitTimeout serverArgs.pollListenURLTimeoutSeconds = defaultPollListenURLTimeout serverArgs.listenAddrHost = defaultListenAddrHost - serverArgs.cockroachLogsDir = baseDir for _, applyOptToArgs := range opts { applyOptToArgs(serverArgs) } + var err error + if serverArgs.baseDir == "" { + baseDir, err := os.MkdirTemp("", "cockroach-testserver") + if err != nil { + return nil, fmt.Errorf("%s: could not create temp directory: %w", testserverMessagePrefix, err) + } + serverArgs.baseDir = baseDir + if serverArgs.cockroachLogsDir == "" { + serverArgs.cockroachLogsDir = baseDir + } + } + log.Printf("cockroach logs directory: %s", serverArgs.cockroachLogsDir) if serverArgs.cockroachBinary != "" { @@ -487,7 +511,7 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) { } mkDir := func(name string) (string, error) { - path := filepath.Join(baseDir, name) + path := filepath.Join(serverArgs.baseDir, name) if err := os.MkdirAll(path, 0755); err != nil { return "", fmt.Errorf("%s: could not create %s directory: %s: %w", testserverMessagePrefix, name, path, err) @@ -629,7 +653,6 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) { serverArgs: *serverArgs, version: v, serverState: stateNew, - baseDir: baseDir, initCmdArgs: initArgs, curTenantID: firstTenantID, nodes: nodes, @@ -677,7 +700,7 @@ func (ts *testServerImpl) StderrForNode(i int) string { // BaseDir returns directory StoreOnDiskOpt writes to if used. func (ts *testServerImpl) BaseDir() string { - return ts.baseDir + return ts.serverArgs.baseDir } // PGURL returns the postgres connection URL to reach the started @@ -878,8 +901,10 @@ func (ts *testServerImpl) Stop() { } ts.mu.RLock() - nodeDir := filepath.Join(ts.baseDir, strconv.Itoa(i)) - if err := os.RemoveAll(nodeDir); err != nil { + nodeDir := filepath.Join(ts.BaseDir(), strconv.Itoa(i)) + if ts.serverArgs.noFileCleanup { + log.Printf("%s: skipping file cleanup of node-dir %s", testserverMessagePrefix, nodeDir) + } else if err := os.RemoveAll(nodeDir); err != nil { log.Printf("error deleting tmp directory %s for node: %s", nodeDir, err) } if closeErr := ts.nodes[i].stdoutBuf.Close(); closeErr != nil { @@ -891,7 +916,11 @@ func (ts *testServerImpl) Stop() { } // Only cleanup on intentional stops. - _ = os.RemoveAll(ts.baseDir) + if ts.serverArgs.noFileCleanup { + log.Printf("%s: skipping file cleanup of base-dir %s", testserverMessagePrefix, ts.BaseDir()) + } else { + _ = os.RemoveAll(ts.BaseDir()) + } } func (ts *testServerImpl) CockroachInit() error { @@ -907,7 +936,7 @@ func (ts *testServerImpl) CockroachInit() error { // Set the working directory of the cockroach process to our temp folder. // This stops cockroach from polluting the project directory with _dump // folders. - ts.initCmd.Dir = ts.baseDir + ts.initCmd.Dir = ts.serverArgs.baseDir err := ts.initCmd.Start() if ts.initCmd.Process != nil { diff --git a/testserver/testserver_test.go b/testserver/testserver_test.go index 01e5e77..7798597 100644 --- a/testserver/testserver_test.go +++ b/testserver/testserver_test.go @@ -234,6 +234,12 @@ func TestRunServer(t *testing.T) { ) }, }, + { + name: "No File Cleanup", + instantiation: func(t *testing.T) (*sql.DB, func()) { + return testserver.NewDBForTest(t, testserver.NoFileCleanup()) + }, + }, } { t.Run(tc.name, func(t *testing.T) { db, stop := tc.instantiation(t) @@ -248,6 +254,35 @@ func TestRunServer(t *testing.T) { } } +func TestBaseDirIsPreservedWhenNoFileCleanupRequested(t *testing.T) { + ts, err := testserver.NewTestServer( + testserver.NoFileCleanup(), + testserver.StoreOnDiskOpt()) + require.NoError(t, err) + baseDir := ts.BaseDir() + + db, err := sql.Open("postgres", ts.PGURL().String()) + require.NoError(t, err) + + _, err = db.Exec("create table store_preservation_test(id int primary key)") + require.NoError(t, err) + _, err = db.Exec("insert into store_preservation_test(id) values (123456789)") + require.NoError(t, err) + ts.Stop() + + newTs, err := testserver.NewTestServer( + testserver.BasedirOverride(baseDir), + testserver.StoreOnDiskOpt()) + require.NoError(t, err) + defer newTs.Stop() + db, err = sql.Open("postgres", newTs.PGURL().String()) + require.NoError(t, err) + row := db.QueryRow("select id from store_preservation_test") + retrivedIntVal := 0 + require.NoError(t, row.Scan(&retrivedIntVal)) + require.Equal(t, 123456789, retrivedIntVal) +} + func TestCockroachBinaryPathOpt(t *testing.T) { crdbBinary := "doesnotexist" _, err := testserver.NewTestServer(testserver.CockroachBinaryPathOpt(crdbBinary)) diff --git a/testserver/testservernode.go b/testserver/testservernode.go index 23668bb..755decd 100644 --- a/testserver/testservernode.go +++ b/testserver/testservernode.go @@ -102,7 +102,7 @@ func (ts *testServerImpl) StartNode(i int) error { // Set the working directory of the cockroach process to our temp folder. // This stops cockroach from polluting the project directory with _dump // folders. - currCmd.Dir = ts.baseDir + currCmd.Dir = ts.BaseDir() if len(ts.nodes[i].stdout) > 0 { wr, err := newFileLogWriter(ts.nodes[i].stdout)