release-23.1: roachtest: clean up artifacts zipping code

pkg/cmd/roachtest/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "test_registry.go",
         "test_runner.go",
         "work_pool.go",
+        "zip_util.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachtest",
     visibility = ["//visibility:private"],
@@ -82,6 +83,7 @@ go_test(
         "main_test.go",
         "test_registry_test.go",
         "test_test.go",
+        "zip_util_test.go",
     ],
     args = ["-test.timeout=55s"],
     data = glob(["testdata/**"]),

pkg/cmd/roachtest/test_runner.go

@@ -11,7 +11,6 @@
 package main
 
 import (
-	"archive/zip"
 	"context"
 	gosql "database/sql"
 	"fmt"
@@ -1683,78 +1682,18 @@ func (we *workerErrors) Err() error {
 	return we.mu.errs[0]
 }
 
-func zipArtifacts(path string) error {
-	f, err := os.Create(filepath.Join(path, "artifacts.zip"))
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-	z := zip.NewWriter(f)
-	rel := func(targetpath string) string {
-		relpath, err := filepath.Rel(path, targetpath)
-		if err != nil {
-			return targetpath
-		}
-		return relpath
-	}
-
-	walk := func(visitor func(string, os.FileInfo) error) error {
-		return filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
-			if err != nil {
-				return err
-			}
-			dir, _ := filepath.Split(rel(path))
-			isTopLevel := dir == ""
-			if !info.IsDir() && isTopLevel && strings.HasSuffix(path, ".zip") {
-				// Skip any top-level zip files, which notably includes itself
-				// and, if present, the debug.zip.
-				return nil
-			}
-			return visitor(path, info)
-		})
-	}
-
-	// Zip all of the files.
-	if err := walk(func(path string, info os.FileInfo) error {
-		if info.IsDir() {
-			return nil
+// zipArtifacts moves everything inside the artifacts dir except any zip files
+// (like debug.zip) into an artifacts.zip file.
+func zipArtifacts(artifactsDir string) error {
+	list, err := filterDirEntries(artifactsDir, func(entry os.DirEntry) bool {
+		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".zip") {
+			// Skip any zip files.
+			return false
 		}
-		w, err := z.Create(rel(path))
-		if err != nil {
-			return err
-		}
-		r, err := os.Open(path)
-		if err != nil {
-			return err
-		}
-		defer r.Close()
-		if _, err := io.Copy(w, r); err != nil {
-			return err
-		}
-		return nil
-	}); err != nil {
-		return err
-	}
-	if err := z.Close(); err != nil {
-		return err
-	}
-	if err := f.Sync(); err != nil {
+		return true
+	})
+	if err != nil {
 		return err
 	}
-
-	// Now that the zip file is there, remove all of the files that went into it.
-	// Note that 'walk' skips the debug.zip and our newly written zip file.
-	root := path
-	return walk(func(path string, info os.FileInfo) error {
-		if path == root {
-			return nil
-		}
-		if err := os.RemoveAll(path); err != nil {
-			return err
-		}
-		if info.IsDir() {
-			return filepath.SkipDir
-		}
-		return nil
-	})
+	return moveToZipArchive("artifacts.zip", artifactsDir, list...)
 }

pkg/cmd/roachtest/zip_util.go

@@ -0,0 +1,98 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package main
+
+import (
+	"archive/zip"
+	"io"
+	"os"
+	"path/filepath"
+	"sort"
+)
+
+// moveToZipArchive creates a zip archive inside rootPath with the files/dirs
+// given as relative paths, after which the given files/dirs are deleted.
+//
+// See filterDirEntries for a convenient way to get a list of files/dirs.
+func moveToZipArchive(archiveName string, rootPath string, relPaths ...string) error {
+	f, err := os.Create(filepath.Join(rootPath, archiveName))
+	if err != nil {
+		return err
+	}
+
+	z := zip.NewWriter(f)
+	for _, relPath := range relPaths {
+		// Walk the given path.
+		if err := filepath.Walk(filepath.Join(rootPath, relPath), func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				return err
+			}
+			if info.IsDir() {
+				// Let Walk recurse inside.
+				return nil
+			}
+			relPath, err := filepath.Rel(rootPath, path)
+			if err != nil {
+				return err
+			}
+			w, err := z.Create(relPath)
+			if err != nil {
+				return err
+			}
+			r, err := os.Open(path)
+			if err != nil {
+				return err
+			}
+			if _, err := io.Copy(w, r); err != nil {
+				_ = r.Close()
+				return err
+			}
+			return r.Close()
+		}); err != nil {
+			_ = f.Close()
+			return err
+		}
+	}
+
+	if err := z.Close(); err != nil {
+		_ = f.Close()
+		return err
+	}
+	if err := f.Close(); err != nil {
+		return err
+	}
+
+	// Now that the zip file is there, remove all the files that went into it.
+	for _, relPath := range relPaths {
+		if err := os.RemoveAll(filepath.Join(rootPath, relPath)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// filterDirEntries lists the given directory, runs a filter function on each
+// entry, and returns the base names of those which passed the filter.
+func filterDirEntries(
+	path string, filter func(entry os.DirEntry) bool,
+) (baseNames []string, _ error) {
+	entries, err := os.ReadDir(path)
+	if err != nil {
+		return nil, err
+	}
+	for _, e := range entries {
+		if filter(e) {
+			baseNames = append(baseNames, e.Name())
+		}
+	}
+	sort.Strings(baseNames)
+	return baseNames, nil
+}

pkg/cmd/roachtest/zip_util_test.go

@@ -0,0 +1,84 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package main
+
+import (
+	"archive/zip"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMoveToZipArchive(t *testing.T) {
+	dir := t.TempDir()
+	p := func(elems ...string) string {
+		return filepath.Join(append([]string{dir}, elems...)...)
+	}
+	require.NoError(t, os.WriteFile(p("a1"), []byte("foo"), 0777))
+	require.NoError(t, os.WriteFile(p("a2"), []byte("foo"), 0777))
+	require.NoError(t, os.WriteFile(p("a3"), []byte("foo"), 0777))
+	require.NoError(t, os.Mkdir(p("dir1"), 0777))
+	require.NoError(t, os.WriteFile(p("dir1", "file1"), []byte("foo"), 0777))
+	require.NoError(t, os.WriteFile(p("dir1", "file2"), []byte("foo"), 0777))
+	require.NoError(t, os.Mkdir(p("dir2"), 0777))
+	require.NoError(t, os.WriteFile(p("dir2", "file1"), []byte("foo"), 0777))
+	require.NoError(t, os.WriteFile(p("dir2", "file2"), []byte("foo"), 0777))
+
+	// expectLs checks the current directory listing of dir.
+	expectLs := func(expected ...string) {
+		t.Helper()
+		var actual []string
+		require.NoError(t, filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
+			require.NoError(t, err)
+			if !info.IsDir() {
+				rel, err := filepath.Rel(dir, path)
+				require.NoError(t, err)
+				actual = append(actual, rel)
+			}
+			return nil
+		}))
+		require.Equal(t, expected, actual)
+	}
+	// expectZip checks the files contained in the given archive; paths must use
+	// slashes.
+	expectZip := func(archiveName string, expected ...string) {
+		r, err := zip.OpenReader(p(archiveName))
+		require.NoError(t, err)
+		var actual []string
+		for _, f := range r.File {
+			actual = append(actual, f.Name)
+		}
+		require.Equal(t, expected, actual)
+		require.NoError(t, r.Close())
+	}
+	j := filepath.Join
+	expectLs("a1", "a2", "a3", j("dir1", "file1"), j("dir1", "file2"), j("dir2", "file1"), j("dir2", "file2"))
+
+	list, err := filterDirEntries(dir, func(entry os.DirEntry) bool {
+		return entry.Name() != "a2" && entry.Name() != "dir2"
+	})
+	require.NoError(t, err)
+	require.Equal(t, []string{"a1", "a3", "dir1"}, list)
+	require.NoError(t, moveToZipArchive("first.zip", dir, list...))
+	expectZip("first.zip", "a1", "a3", "dir1/file1", "dir1/file2")
+	expectLs("a2", j("dir2", "file1"), j("dir2", "file2"), "first.zip")
+
+	list, err = filterDirEntries(dir, func(entry os.DirEntry) bool {
+		return !strings.HasSuffix(entry.Name(), ".zip")
+	})
+	require.NoError(t, err)
+	require.NoError(t, moveToZipArchive("second.zip", dir, list...))
+	expectZip("second.zip", "a2", "dir2/file1", "dir2/file2")
+	expectLs("first.zip", "second.zip")
+}