From c9f28e5a800c6e393486bd0cf9f40911cb746a9f Mon Sep 17 00:00:00 2001
From: Tobias Grieger <tobias.b.grieger@gmail.com>
Date: Fri, 28 May 2021 11:45:29 +0200
Subject: [PATCH 1/2] roachtest: don't duplicate log files

Since we want to have the logs for down nodes as well, roachtest copies
the log directories to artifacts. It is therefore not necessary to
include the logs in the `debug.zip` in addition to that.

Release note: None
---
 pkg/cmd/roachtest/cluster.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/pkg/cmd/roachtest/cluster.go b/pkg/cmd/roachtest/cluster.go
index 87fedef2b34b..9ca329efeea3 100644
--- a/pkg/cmd/roachtest/cluster.go
+++ b/pkg/cmd/roachtest/cluster.go
@@ -1698,9 +1698,12 @@ func (c *cluster) FetchDebugZip(ctx context.Context) error {
 		// waste our time.
 		for i := 1; i <= c.spec.NodeCount; i++ {
 			// `./cockroach debug zip` is noisy. Suppress the output unless it fails.
+			//
+			// Ignore the files in the the log directory; we pull the logs separately anyway
+			// so this would only cause duplication.
 			si := strconv.Itoa(i)
 			output, err := execCmdWithBuffer(ctx, c.l, roachprod, "ssh", c.name+":"+si, "--",
-				"./cockroach", "debug", "zip", "--url", "{pgurl:"+si+"}", zipName)
+				"./cockroach", "debug", "zip", "--exclude-files='*.log,*.txt,*.pprof'", "--url", "{pgurl:"+si+"}", zipName)
 			if err != nil {
 				c.l.Printf("./cockroach debug zip failed: %s", output)
 				if i < c.spec.NodeCount {

From 7ad39e6af943a9d4deca9ee767d786af0f23163b Mon Sep 17 00:00:00 2001
From: Tobias Grieger <tobias.b.grieger@gmail.com>
Date: Fri, 28 May 2021 12:28:05 +0200
Subject: [PATCH 2/2] roachtest: zip the artifacts for each test on teamcity

Often you want to download the artifacts for a specific test,
but TeamCity doesn't let you do that. Zip the artifacts for
each test as a workaround.

TeamCity is good at navigating into zip files, so we don't lose
anything in terms of poking at a test directly from TeamCity.

cc @cockroachdb/test-eng

Release note: None
---
 pkg/cmd/roachtest/test_runner.go | 83 ++++++++++++++++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/pkg/cmd/roachtest/test_runner.go b/pkg/cmd/roachtest/test_runner.go
index b7e38e9a4503..d0640b63ceb2 100644
--- a/pkg/cmd/roachtest/test_runner.go
+++ b/pkg/cmd/roachtest/test_runner.go
@@ -11,6 +11,7 @@
 package main
 
 import (
+	"archive/zip"
 	"context"
 	"fmt"
 	"html"
@@ -710,6 +711,14 @@ caffeinate ./roachstress.sh %s
 		if teamCity {
 			shout(ctx, l, stdout, "##teamcity[testFinished name='%s' flowId='%s']", t.Name(), t.Name())
 
+			// Zip the artifacts. This improves the TeamCity UX where we can navigate
+			// through zip files just fine, but we can't download subtrees of the
+			// artifacts storage. By zipping we get this capability as we can just
+			// download the zip file for the failing test instead.
+			if err := zipArtifacts(t.artifactsDir); err != nil {
+				l.Printf("unable to zip artifacts: %s", err)
+			}
+
 			if t.artifactsSpec != "" {
 				// Tell TeamCity to collect this test's artifacts now. The TC job
 				// also collects the artifacts directory wholesale at the end, but
@@ -1254,3 +1263,77 @@ func (we *workerErrors) Err() error {
 	// 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
+			}
+			if !info.IsDir() && 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
+		}
+		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 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
+	})
+}