diff --git a/pkg/cmd/roachtest/BUILD.bazel b/pkg/cmd/roachtest/BUILD.bazel index 588adbd2482c..5d19c0f44e02 100644 --- a/pkg/cmd/roachtest/BUILD.bazel +++ b/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"], @@ -81,6 +82,7 @@ go_test( "main_test.go", "test_registry_test.go", "test_test.go", + "zip_util_test.go", ], args = ["-test.timeout=55s"], data = glob(["testdata/**"]), diff --git a/pkg/cmd/roachtest/test_runner.go b/pkg/cmd/roachtest/test_runner.go index bcf1fc4c936d..00ad8597e4fa 100644 --- a/pkg/cmd/roachtest/test_runner.go +++ b/pkg/cmd/roachtest/test_runner.go @@ -11,7 +11,6 @@ package main import ( - "archive/zip" "context" gosql "database/sql" "fmt" @@ -1678,78 +1677,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...) } diff --git a/pkg/cmd/roachtest/zip_util.go b/pkg/cmd/roachtest/zip_util.go new file mode 100644 index 000000000000..60c2b702c8d6 --- /dev/null +++ b/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 +} diff --git a/pkg/cmd/roachtest/zip_util_test.go b/pkg/cmd/roachtest/zip_util_test.go new file mode 100644 index 000000000000..06ef259c43db --- /dev/null +++ b/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") +}