refactor(nixery): Extract layering logic into separate package r/4105

This will be required for making a standalone, Nixery-style image
builder function usable from Nix.

Change-Id: I5e36348bd4c32d249d56f6628cd046916691319f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/5601
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

3 files changed, 8 insertions, 358 deletions

diff --git a/tools/nixery/builder/archive.go b/tools/nixery/builder/archive.go
index 3bc02ab4d5b8..8763e4cb8566 100644
--- a/tools/nixery/builder/archive.go
+++ b/tools/nixery/builder/archive.go
@@ -16,6 +16,8 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+
+	"github.com/google/nixery/layers"
 )
 
 // Create a new compressed tarball from each of the paths in the list
@@ -23,7 +25,7 @@ import (
 //
 // The uncompressed tarball is hashed because image manifests must
 // contain both the hashes of compressed and uncompressed layers.
-func packStorePaths(l *layer, w io.Writer) (string, error) {
+func packStorePaths(l *layers.Layer, w io.Writer) (string, error) {
 	shasum := sha256.New()
 	gz := gzip.NewWriter(w)
 	multi := io.MultiWriter(shasum, gz)
diff --git a/tools/nixery/builder/builder.go b/tools/nixery/builder/builder.go
index 37c9b9fcb763..901373b57e99 100644
--- a/tools/nixery/builder/builder.go
+++ b/tools/nixery/builder/builder.go
@@ -23,6 +23,7 @@ import (
 	"strings"
 
 	"github.com/google/nixery/config"
+	"github.com/google/nixery/layers"
 	"github.com/google/nixery/manifest"
 	"github.com/google/nixery/storage"
 	log "github.com/sirupsen/logrus"
@@ -39,7 +40,7 @@ type State struct {
 	Storage storage.Backend
 	Cache   *LocalCache
 	Cfg     config.Config
-	Pop     Popularity
+	Pop     layers.Popularity
 }
 
 // Architecture represents the possible CPU architectures for which
@@ -117,7 +118,7 @@ type ImageResult struct {
 	Pkgs  []string `json:"pkgs"`
 
 	// These fields are populated in case of success
-	Graph        runtimeGraph `json:"runtimeGraph"`
+	Graph        layers.RuntimeGraph `json:"runtimeGraph"`
 	SymlinkLayer struct {
 		Size    int    `json:"size"`
 		TarHash string `json:"tarHash"`
@@ -281,7 +282,7 @@ func prepareImage(s *State, image *Image) (*ImageResult, error) {
 // added only after successful uploads, which guarantees that entries
 // retrieved from the cache are present in the bucket.
 func prepareLayers(ctx context.Context, s *State, image *Image, result *ImageResult) ([]manifest.Entry, error) {
-	grouped := groupLayers(&result.Graph, &s.Pop, LayerBudget)
+	grouped := layers.GroupLayers(&result.Graph, &s.Pop, LayerBudget)
 
 	var entries []manifest.Entry
 
@@ -318,7 +319,7 @@ func prepareLayers(ctx context.Context, s *State, image *Image, result *ImageRes
 
 			var pkgs []string
 			for _, p := range l.Contents {
-				pkgs = append(pkgs, packageFromPath(p))
+				pkgs = append(pkgs, layers.PackageFromPath(p))
 			}
 
 			log.WithFields(log.Fields{
diff --git a/tools/nixery/builder/layers.go b/tools/nixery/builder/layers.go
deleted file mode 100644
index 5e37e626810f..000000000000
--- a/tools/nixery/builder/layers.go
+++ /dev/null
@@ -1,353 +0,0 @@
-// Copyright 2022 The TVL Contributors
-// SPDX-License-Identifier: Apache-2.0
-
-// This package reads an export reference graph (i.e. a graph representing the
-// runtime dependencies of a set of derivations) created by Nix and groups it in
-// a way that is likely to match the grouping for other derivation sets with
-// overlapping dependencies.
-//
-// This is used to determine which derivations to include in which layers of a
-// container image.
-//
-// # Inputs
-//
-// * a graph of Nix runtime dependencies, generated via exportReferenceGraph
-// * popularity values of each package in the Nix package set (in the form of a
-//   direct reference count)
-// * a maximum number of layers to allocate for the image (the "layer budget")
-//
-// # Algorithm
-//
-// It works by first creating a (directed) dependency tree:
-//
-// img (root node)
-// │
-// ├───> A ─────┐
-// │            v
-// ├───> B ───> E
-// │            ^
-// ├───> C ─────┘
-// │     │
-// │     v
-// └───> D ───> F
-//       │
-//       └────> G
-//
-// Each node (i.e. package) is then visited to determine how important
-// it is to separate this node into its own layer, specifically:
-//
-// 1. Is the node within a certain threshold percentile of absolute
-//    popularity within all of nixpkgs? (e.g. `glibc`, `openssl`)
-//
-// 2. Is the node's runtime closure above a threshold size? (e.g. 100MB)
-//
-// In either case, a bit is flipped for this node representing each
-// condition and an edge to it is inserted directly from the image
-// root, if it does not already exist.
-//
-// For the rest of the example we assume 'G' is above the threshold
-// size and 'E' is popular.
-//
-// This tree is then transformed into a dominator tree:
-//
-// img
-// │
-// ├───> A
-// ├───> B
-// ├───> C
-// ├───> E
-// ├───> D ───> F
-// └───> G
-//
-// Specifically this means that the paths to A, B, C, E, G, and D
-// always pass through the root (i.e. are dominated by it), whilst F
-// is dominated by D (all paths go through it).
-//
-// The top-level subtrees are considered as the initially selected
-// layers.
-//
-// If the list of layers fits within the layer budget, it is returned.
-//
-// Otherwise, a merge rating is calculated for each layer. This is the
-// product of the layer's total size and its root node's popularity.
-//
-// Layers are then merged in ascending order of merge ratings until
-// they fit into the layer budget.
-//
-// # Threshold values
-//
-// Threshold values for the partitioning conditions mentioned above
-// have not yet been determined, but we will make a good first guess
-// based on gut feeling and proceed to measure their impact on cache
-// hits/misses.
-//
-// # Example
-//
-// Using the logic described above as well as the example presented in
-// the introduction, this program would create the following layer
-// groupings (assuming no additional partitioning):
-//
-// Layer budget: 1
-// Layers: { A, B, C, D, E, F, G }
-//
-// Layer budget: 2
-// Layers: { G }, { A, B, C, D, E, F }
-//
-// Layer budget: 3
-// Layers: { G }, { E }, { A, B, C, D, F }
-//
-// Layer budget: 4
-// Layers: { G }, { E }, { D, F }, { A, B, C }
-//
-// ...
-//
-// Layer budget: 10
-// Layers: { E }, { D, F }, { A }, { B }, { C }
-package builder
-
-import (
-	"crypto/sha1"
-	"fmt"
-	"regexp"
-	"sort"
-	"strings"
-
-	log "github.com/sirupsen/logrus"
-	"gonum.org/v1/gonum/graph/flow"
-	"gonum.org/v1/gonum/graph/simple"
-)
-
-// runtimeGraph represents structured information from Nix about the runtime
-// dependencies of a derivation.
-//
-// This is generated in Nix by using the exportReferencesGraph feature.
-type runtimeGraph struct {
-	References struct {
-		Graph []string `json:"graph"`
-	} `json:"exportReferencesGraph"`
-
-	Graph []struct {
-		Size uint64   `json:"closureSize"`
-		Path string   `json:"path"`
-		Refs []string `json:"references"`
-	} `json:"graph"`
-}
-
-// Popularity data for each Nix package that was calculated in advance.
-//
-// Popularity is a number from 1-100 that represents the
-// popularity percentile in which this package resides inside
-// of the nixpkgs tree.
-type Popularity = map[string]int
-
-// Layer represents the data returned for each layer that Nix should
-// build for the container image.
-type layer struct {
-	Contents    []string `json:"contents"`
-	MergeRating uint64
-}
-
-// Hash the contents of a layer to create a deterministic identifier that can be
-// used for caching.
-func (l *layer) Hash() string {
-	sum := sha1.Sum([]byte(strings.Join(l.Contents, ":")))
-	return fmt.Sprintf("%x", sum)
-}
-
-func (a layer) merge(b layer) layer {
-	a.Contents = append(a.Contents, b.Contents...)
-	a.MergeRating += b.MergeRating
-	return a
-}
-
-// closure as pointed to by the graph nodes.
-type closure struct {
-	GraphID    int64
-	Path       string
-	Size       uint64
-	Refs       []string
-	Popularity int
-}
-
-func (c *closure) ID() int64 {
-	return c.GraphID
-}
-
-var nixRegexp = regexp.MustCompile(`^/nix/store/[a-z0-9]+-`)
-
-// PackageFromPath returns the name of a Nix package based on its
-// output store path.
-func packageFromPath(path string) string {
-	return nixRegexp.ReplaceAllString(path, "")
-}
-
-// DOTID provides a human-readable package name. The name stems from
-// the dot format used by GraphViz, into which the dependency graph
-// can be rendered.
-func (c *closure) DOTID() string {
-	return packageFromPath(c.Path)
-}
-
-// bigOrPopular checks whether this closure should be considered for
-// separation into its own layer, even if it would otherwise only
-// appear in a subtree of the dominator tree.
-func (c *closure) bigOrPopular() bool {
-	const sizeThreshold = 100 * 1000000 // 100MB
-
-	if c.Size > sizeThreshold {
-		return true
-	}
-
-	// Threshold value is picked arbitrarily right now. The reason
-	// for this is that some packages (such as `cacert`) have very
-	// few direct dependencies, but are required by pretty much
-	// everything.
-	if c.Popularity >= 100 {
-		return true
-	}
-
-	return false
-}
-
-func insertEdges(graph *simple.DirectedGraph, cmap *map[string]*closure, node *closure) {
-	// Big or popular nodes get a separate edge from the top to
-	// flag them for their own layer.
-	if node.bigOrPopular() && !graph.HasEdgeFromTo(0, node.ID()) {
-		edge := graph.NewEdge(graph.Node(0), node)
-		graph.SetEdge(edge)
-	}
-
-	for _, c := range node.Refs {
-		// Nix adds a self reference to each node, which
-		// should not be inserted.
-		if c != node.Path {
-			edge := graph.NewEdge(node, (*cmap)[c])
-			graph.SetEdge(edge)
-		}
-	}
-}
-
-// Create a graph structure from the references supplied by Nix.
-func buildGraph(refs *runtimeGraph, pop *Popularity) *simple.DirectedGraph {
-	cmap := make(map[string]*closure)
-	graph := simple.NewDirectedGraph()
-
-	// Insert all closures into the graph, as well as a fake root
-	// closure which serves as the top of the tree.
-	//
-	// A map from store paths to IDs is kept to actually insert
-	// edges below.
-	root := &closure{
-		GraphID: 0,
-		Path:    "image_root",
-	}
-	graph.AddNode(root)
-
-	for idx, c := range refs.Graph {
-		node := &closure{
-			GraphID: int64(idx + 1), // inc because of root node
-			Path:    c.Path,
-			Size:    c.Size,
-			Refs:    c.Refs,
-		}
-
-		// The packages `nss-cacert` and `iana-etc` are added
-		// by Nixery to *every single image* and should have a
-		// very high popularity.
-		//
-		// Other popularity values are populated from the data
-		// set assembled by Nixery's popcount.
-		id := node.DOTID()
-		if strings.HasPrefix(id, "nss-cacert") || strings.HasPrefix(id, "iana-etc") {
-			// glibc has ~300k references, these packages need *more*
-			node.Popularity = 500000
-		} else if p, ok := (*pop)[id]; ok {
-			node.Popularity = p
-		} else {
-			node.Popularity = 1
-		}
-
-		graph.AddNode(node)
-		cmap[c.Path] = node
-	}
-
-	// Insert the top-level closures with edges from the root
-	// node, then insert all edges for each closure.
-	for _, p := range refs.References.Graph {
-		edge := graph.NewEdge(root, cmap[p])
-		graph.SetEdge(edge)
-	}
-
-	for _, c := range cmap {
-		insertEdges(graph, &cmap, c)
-	}
-
-	return graph
-}
-
-// Extracts a subgraph starting at the specified root from the
-// dominator tree. The subgraph is converted into a flat list of
-// layers, each containing the store paths and merge rating.
-func groupLayer(dt *flow.DominatorTree, root *closure) layer {
-	size := root.Size
-	contents := []string{root.Path}
-	children := dt.DominatedBy(root.ID())
-
-	// This iteration does not use 'range' because the list being
-	// iterated is modified during the iteration (yes, I'm sorry).
-	for i := 0; i < len(children); i++ {
-		child := children[i].(*closure)
-		size += child.Size
-		contents = append(contents, child.Path)
-		children = append(children, dt.DominatedBy(child.ID())...)
-	}
-
-	// Contents are sorted to ensure that hashing is consistent
-	sort.Strings(contents)
-
-	return layer{
-		Contents:    contents,
-		MergeRating: uint64(root.Popularity) * size,
-	}
-}
-
-// Calculate the dominator tree of the entire package set and group
-// each top-level subtree into a layer.
-//
-// Layers are merged together until they fit into the layer budget,
-// based on their merge rating.
-func dominate(budget int, graph *simple.DirectedGraph) []layer {
-	dt := flow.Dominators(graph.Node(0), graph)
-
-	var layers []layer
-	for _, n := range dt.DominatedBy(dt.Root().ID()) {
-		layers = append(layers, groupLayer(&dt, n.(*closure)))
-	}
-
-	sort.Slice(layers, func(i, j int) bool {
-		return layers[i].MergeRating < layers[j].MergeRating
-	})
-
-	if len(layers) > budget {
-		log.WithFields(log.Fields{
-			"layers": len(layers),
-			"budget": budget,
-		}).Info("ideal image exceeds layer budget")
-	}
-
-	for len(layers) > budget {
-		merged := layers[0].merge(layers[1])
-		layers[1] = merged
-		layers = layers[1:]
-	}
-
-	return layers
-}
-
-// groupLayers applies the algorithm described above the its input and returns a
-// list of layers, each consisting of a list of Nix store paths that it should
-// contain.
-func groupLayers(refs *runtimeGraph, pop *Popularity, budget int) []layer {
-	graph := buildGraph(refs, pop)
-	return dominate(budget, graph)
-}