2 files changed, 336 insertions, 0 deletions
diff --git a/src/internal/coverage/pods/pods.go b/src/internal/coverage/pods/pods.go
new file mode 100644
index 0000000..432c7b6
--- /dev/null
+++ b/src/internal/coverage/pods/pods.go
@@ -0,0 +1,194 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package pods
+
+import (
+	"fmt"
+	"internal/coverage"
+	"os"
+	"path/filepath"
+	"regexp"
+	"sort"
+	"strconv"
+)
+
+// Pod encapsulates a set of files emitted during the executions of a
+// coverage-instrumented binary. Each pod contains a single meta-data
+// file, and then 0 or more counter data files that refer to that
+// meta-data file. Pods are intended to simplify processing of
+// coverage output files in the case where we have several coverage
+// output directories containing output files derived from more
+// than one instrumented executable. In the case where the files that
+// make up a pod are spread out across multiple directories, each
+// element of the "Origins" field below will be populated with the
+// index of the originating directory for the corresponding counter
+// data file (within the slice of input dirs handed to CollectPods).
+// The ProcessIDs field will be populated with the process ID of each
+// data file in the CounterDataFiles slice.
+type Pod struct {
+	MetaFile         string
+	CounterDataFiles []string
+	Origins          []int
+	ProcessIDs       []int
+}
+
+// CollectPods visits the files contained within the directories in
+// the list 'dirs', collects any coverage-related files, partitions
+// them into pods, and returns a list of the pods to the caller, along
+// with an error if something went wrong during directory/file
+// reading.
+//
+// CollectPods skips over any file that is not related to coverage
+// (e.g. avoids looking at things that are not meta-data files or
+// counter-data files). CollectPods also skips over 'orphaned' counter
+// data files (e.g. counter data files for which we can't find the
+// corresponding meta-data file). If "warn" is true, CollectPods will
+// issue warnings to stderr when it encounters non-fatal problems (for
+// orphans or a directory with no meta-data files).
+func CollectPods(dirs []string, warn bool) ([]Pod, error) {
+	files := []string{}
+	dirIndices := []int{}
+	for k, dir := range dirs {
+		dents, err := os.ReadDir(dir)
+		if err != nil {
+			return nil, err
+		}
+		for _, e := range dents {
+			if e.IsDir() {
+				continue
+			}
+			files = append(files, filepath.Join(dir, e.Name()))
+			dirIndices = append(dirIndices, k)
+		}
+	}
+	return collectPodsImpl(files, dirIndices, warn), nil
+}
+
+// CollectPodsFromFiles functions the same as "CollectPods" but
+// operates on an explicit list of files instead of a directory.
+func CollectPodsFromFiles(files []string, warn bool) []Pod {
+	return collectPodsImpl(files, nil, warn)
+}
+
+type fileWithAnnotations struct {
+	file   string
+	origin int
+	pid    int
+}
+
+type protoPod struct {
+	mf       string
+	elements []fileWithAnnotations
+}
+
+// collectPodsImpl examines the specified list of files and picks out
+// subsets that correspond to coverage pods. The first stage in this
+// process is collecting a set { M1, M2, ... MN } where each M_k is a
+// distinct coverage meta-data file. We then create a single pod for
+// each meta-data file M_k, then find all of the counter data files
+// that refer to that meta-data file (recall that the counter data
+// file name incorporates the meta-data hash), and add the counter
+// data file to the appropriate pod.
+//
+// This process is complicated by the fact that we need to keep track
+// of directory indices for counter data files. Here is an example to
+// motivate:
+//
+//	directory 1:
+//
+// M1   covmeta.9bbf1777f47b3fcacb05c38b035512d6
+// C1   covcounters.9bbf1777f47b3fcacb05c38b035512d6.1677673.1662138360208416486
+// C2   covcounters.9bbf1777f47b3fcacb05c38b035512d6.1677637.1662138359974441782
+//
+//	directory 2:
+//
+// M2   covmeta.9bbf1777f47b3fcacb05c38b035512d6
+// C3   covcounters.9bbf1777f47b3fcacb05c38b035512d6.1677445.1662138360208416480
+// C4   covcounters.9bbf1777f47b3fcacb05c38b035512d6.1677677.1662138359974441781
+// M3   covmeta.a723844208cea2ae80c63482c78b2245
+// C5   covcounters.a723844208cea2ae80c63482c78b2245.3677445.1662138360208416480
+// C6   covcounters.a723844208cea2ae80c63482c78b2245.1877677.1662138359974441781
+//
+// In these two directories we have three meta-data files, but only
+// two are distinct, meaning that we'll wind up with two pods. The
+// first pod (with meta-file M1) will have four counter data files
+// (C1, C2, C3, C4) and the second pod will have two counter data files
+// (C5, C6).
+func collectPodsImpl(files []string, dirIndices []int, warn bool) []Pod {
+	metaRE := regexp.MustCompile(fmt.Sprintf(`^%s\.(\S+)$`, coverage.MetaFilePref))
+	mm := make(map[string]protoPod)
+	for _, f := range files {
+		base := filepath.Base(f)
+		if m := metaRE.FindStringSubmatch(base); m != nil {
+			tag := m[1]
+			// We need to allow for the possibility of duplicate
+			// meta-data files. If we hit this case, use the
+			// first encountered as the canonical version.
+			if _, ok := mm[tag]; !ok {
+				mm[tag] = protoPod{mf: f}
+			}
+			// FIXME: should probably check file length and hash here for
+			// the duplicate.
+		}
+	}
+	counterRE := regexp.MustCompile(fmt.Sprintf(coverage.CounterFileRegexp, coverage.CounterFilePref))
+	for k, f := range files {
+		base := filepath.Base(f)
+		if m := counterRE.FindStringSubmatch(base); m != nil {
+			tag := m[1] // meta hash
+			pid, err := strconv.Atoi(m[2])
+			if err != nil {
+				continue
+			}
+			if v, ok := mm[tag]; ok {
+				idx := -1
+				if dirIndices != nil {
+					idx = dirIndices[k]
+				}
+				fo := fileWithAnnotations{file: f, origin: idx, pid: pid}
+				v.elements = append(v.elements, fo)
+				mm[tag] = v
+			} else {
+				if warn {
+					warning("skipping orphaned counter file: %s", f)
+				}
+			}
+		}
+	}
+	if len(mm) == 0 {
+		if warn {
+			warning("no coverage data files found")
+		}
+		return nil
+	}
+	pods := make([]Pod, 0, len(mm))
+	for _, p := range mm {
+		sort.Slice(p.elements, func(i, j int) bool {
+			return p.elements[i].file < p.elements[j].file
+		})
+		pod := Pod{
+			MetaFile:         p.mf,
+			CounterDataFiles: make([]string, 0, len(p.elements)),
+			Origins:          make([]int, 0, len(p.elements)),
+			ProcessIDs:       make([]int, 0, len(p.elements)),
+		}
+		for _, e := range p.elements {
+			pod.CounterDataFiles = append(pod.CounterDataFiles, e.file)
+			pod.Origins = append(pod.Origins, e.origin)
+			pod.ProcessIDs = append(pod.ProcessIDs, e.pid)
+		}
+		pods = append(pods, pod)
+	}
+	sort.Slice(pods, func(i, j int) bool {
+		return pods[i].MetaFile < pods[j].MetaFile
+	})
+	return pods
+}
+
+func warning(s string, a ...interface{}) {
+	fmt.Fprintf(os.Stderr, "warning: ")
+	fmt.Fprintf(os.Stderr, s, a...)
+	fmt.Fprintf(os.Stderr, "\n")
+}
diff --git a/src/internal/coverage/pods/pods_test.go b/src/internal/coverage/pods/pods_test.go
new file mode 100644
index 0000000..28914c5
--- /dev/null
+++ b/src/internal/coverage/pods/pods_test.go
@@ -0,0 +1,142 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package pods_test
+
+import (
+	"crypto/md5"
+	"fmt"
+	"internal/coverage"
+	"internal/coverage/pods"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+)
+
+func TestPodCollection(t *testing.T) {
+	//testenv.MustHaveGoBuild(t)
+
+	mkdir := func(d string, perm os.FileMode) string {
+		dp := filepath.Join(t.TempDir(), d)
+		if err := os.Mkdir(dp, perm); err != nil {
+			t.Fatal(err)
+		}
+		return dp
+	}
+
+	mkfile := func(d string, fn string) string {
+		fp := filepath.Join(d, fn)
+		if err := os.WriteFile(fp, []byte("foo"), 0666); err != nil {
+			t.Fatal(err)
+		}
+		return fp
+	}
+
+	mkmeta := func(dir string, tag string) string {
+		hash := md5.Sum([]byte(tag))
+		fn := fmt.Sprintf("%s.%x", coverage.MetaFilePref, hash)
+		return mkfile(dir, fn)
+	}
+
+	mkcounter := func(dir string, tag string, nt int) string {
+		hash := md5.Sum([]byte(tag))
+		dummyPid := int(42)
+		fn := fmt.Sprintf(coverage.CounterFileTempl, coverage.CounterFilePref, hash, dummyPid, nt)
+		return mkfile(dir, fn)
+	}
+
+	trim := func(path string) string {
+		b := filepath.Base(path)
+		d := filepath.Dir(path)
+		db := filepath.Base(d)
+		return db + "/" + b
+	}
+
+	podToString := func(p pods.Pod) string {
+		rv := trim(p.MetaFile) + " [\n"
+		for k, df := range p.CounterDataFiles {
+			rv += trim(df)
+			if p.Origins != nil {
+				rv += fmt.Sprintf(" o:%d", p.Origins[k])
+			}
+			rv += "\n"
+		}
+		return rv + "]"
+	}
+
+	// Create a couple of directories.
+	o1 := mkdir("o1", 0777)
+	o2 := mkdir("o2", 0777)
+
+	// Add some random files (not coverage related)
+	mkfile(o1, "blah.txt")
+	mkfile(o1, "something.exe")
+
+	// Add a meta-data file with two counter files to first dir.
+	mkmeta(o1, "m1")
+	mkcounter(o1, "m1", 1)
+	mkcounter(o1, "m1", 2)
+	mkcounter(o1, "m1", 2)
+
+	// Add a counter file with no associated meta file.
+	mkcounter(o1, "orphan", 9)
+
+	// Add a meta-data file with three counter files to second dir.
+	mkmeta(o2, "m2")
+	mkcounter(o2, "m2", 1)
+	mkcounter(o2, "m2", 2)
+	mkcounter(o2, "m2", 3)
+
+	// Add a duplicate of the first meta-file and a corresponding
+	// counter file to the second dir. This is intended to capture
+	// the scenario where we have two different runs of the same
+	// coverage-instrumented binary, but with the output files
+	// sent to separate directories.
+	mkmeta(o2, "m1")
+	mkcounter(o2, "m1", 11)
+
+	// Collect pods.
+	podlist, err := pods.CollectPods([]string{o1, o2}, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify pods
+	if len(podlist) != 2 {
+		t.Fatalf("expected 2 pods got %d pods", len(podlist))
+	}
+
+	for k, p := range podlist {
+		t.Logf("%d: mf=%s\n", k, p.MetaFile)
+	}
+
+	expected := []string{
+		`o1/covmeta.ae7be26cdaa742ca148068d5ac90eaca [
+o1/covcounters.ae7be26cdaa742ca148068d5ac90eaca.42.1 o:0
+o1/covcounters.ae7be26cdaa742ca148068d5ac90eaca.42.2 o:0
+o2/covcounters.ae7be26cdaa742ca148068d5ac90eaca.42.11 o:1
+]`,
+		`o2/covmeta.aaf2f89992379705dac844c0a2a1d45f [
+o2/covcounters.aaf2f89992379705dac844c0a2a1d45f.42.1 o:1
+o2/covcounters.aaf2f89992379705dac844c0a2a1d45f.42.2 o:1
+o2/covcounters.aaf2f89992379705dac844c0a2a1d45f.42.3 o:1
+]`,
+	}
+	for k, exp := range expected {
+		got := podToString(podlist[k])
+		if exp != got {
+			t.Errorf("pod %d: expected:\n%s\ngot:\n%s", k, exp, got)
+		}
+	}
+
+	// Check handling of bad/unreadable dir.
+	if runtime.GOOS == "linux" {
+		dbad := "/dev/null"
+		_, err = pods.CollectPods([]string{dbad}, true)
+		if err == nil {
+			t.Errorf("exected error due to unreadable dir")
+		}
+	}
+}