core/pkgs/build-support/references-by-popularity/closure-graph.py
2024-05-13 11:34:52 -04:00

568 lines
16 KiB
Python

# IMPORTANT: Making changes?
#
# Validate your changes with python3 ./closure-graph.py --test
# Using a simple algorithm, convert the references to a path in to a
# sorted list of dependent paths based on how often they're referenced
# and how deep in the tree they live. Equally-"popular" paths are then
# sorted by name.
#
# The existing writeClosure prints the paths in a simple ascii-based
# sorting of the paths.
#
# Sorting the paths by graph improves the chances that the difference
# between two builds appear near the end of the list, instead of near
# the beginning. This makes a difference for Nix builds which export a
# closure for another program to consume, if that program implements its
# own level of binary diffing.
#
# For an example, Docker Images. If each store path is a separate layer
# then Docker Images can be very efficiently transfered between systems,
# and we get very good cache reuse between images built with the same
# version of Nixpkgs. However, since Docker only reliably supports a
# small number of layers (42) it is important to pick the individual
# layers carefully. By storing very popular store paths in the first 40
# layers, we improve the chances that the next Docker image will share
# many of those layers.*
#
# Given the dependency tree:
#
# A - B - C - D -\
# \ \ \ \
# \ \ \ \
# \ \ - E ---- F
# \- G
#
# Nodes which have multiple references are duplicated:
#
# A - B - C - D - F
# \ \ \
# \ \ \- E - F
# \ \
# \ \- E - F
# \
# \- G
#
# Each leaf node is now replaced by a counter defaulted to 1:
#
# A - B - C - D - (F:1)
# \ \ \
# \ \ \- E - (F:1)
# \ \
# \ \- E - (F:1)
# \
# \- (G:1)
#
# Then each leaf counter is merged with its parent node, replacing the
# parent node with a counter of 1, and each existing counter being
# incremented by 1. That is to say `- D - (F:1)` becomes `- (D:1, F:2)`:
#
# A - B - C - (D:1, F:2)
# \ \ \
# \ \ \- (E:1, F:2)
# \ \
# \ \- (E:1, F:2)
# \
# \- (G:1)
#
# Then each leaf counter is merged with its parent node again, merging
# any counters, then incrementing each:
#
# A - B - (C:1, D:2, E:2, F:5)
# \ \
# \ \- (E:1, F:2)
# \
# \- (G:1)
#
# And again:
#
# A - (B:1, C:2, D:3, E:4, F:8)
# \
# \- (G:1)
#
# And again:
#
# (A:1, B:2, C:3, D:4, E:5, F:9, G:2)
#
# and then paths have the following "popularity":
#
# A 1
# B 2
# C 3
# D 4
# E 5
# F 9
# G 2
#
# and the popularity contest would result in the paths being printed as:
#
# F
# E
# D
# C
# B
# G
# A
#
# * Note: People who have used a Dockerfile before assume Docker's
# Layers are inherently ordered. However, this is not true -- Docker
# layers are content-addressable and are not explicitly layered until
# they are composed in to an Image.
import sys
import json
import unittest
from pprint import pprint
from collections import defaultdict
def debug(msg, *args, **kwargs):
if False:
print(
"DEBUG: {}".format(
msg.format(*args, **kwargs)
),
file=sys.stderr
)
# Find paths in the original dataset which are never referenced by
# any other paths
def find_roots(closures):
roots = [];
for closure in closures:
path = closure['path']
if not any_refer_to(path, closures):
roots.append(path)
return roots
class TestFindRoots(unittest.TestCase):
def test_find_roots(self):
self.assertCountEqual(
find_roots([
{
"path": "/nix/store/foo",
"references": [
"/nix/store/foo",
"/nix/store/bar"
]
},
{
"path": "/nix/store/bar",
"references": [
"/nix/store/bar",
"/nix/store/tux"
]
},
{
"path": "/nix/store/hello",
"references": [
]
}
]),
["/nix/store/foo", "/nix/store/hello"]
)
def any_refer_to(path, closures):
for closure in closures:
if path != closure['path']:
if path in closure['references']:
return True
return False
class TestAnyReferTo(unittest.TestCase):
def test_has_references(self):
self.assertTrue(
any_refer_to(
"/nix/store/bar",
[
{
"path": "/nix/store/foo",
"references": [
"/nix/store/bar"
]
},
]
),
)
def test_no_references(self):
self.assertFalse(
any_refer_to(
"/nix/store/foo",
[
{
"path": "/nix/store/foo",
"references": [
"/nix/store/foo",
"/nix/store/bar"
]
},
]
),
)
def all_paths(closures):
paths = []
for closure in closures:
paths.append(closure['path'])
paths.extend(closure['references'])
paths.sort()
return list(set(paths))
class TestAllPaths(unittest.TestCase):
def test_returns_all_paths(self):
self.assertCountEqual(
all_paths([
{
"path": "/nix/store/foo",
"references": [
"/nix/store/foo",
"/nix/store/bar"
]
},
{
"path": "/nix/store/bar",
"references": [
"/nix/store/bar",
"/nix/store/tux"
]
},
{
"path": "/nix/store/hello",
"references": [
]
}
]),
["/nix/store/foo", "/nix/store/bar", "/nix/store/hello", "/nix/store/tux",]
)
def test_no_references(self):
self.assertFalse(
any_refer_to(
"/nix/store/foo",
[
{
"path": "/nix/store/foo",
"references": [
"/nix/store/foo",
"/nix/store/bar"
]
},
]
),
)
# Convert:
#
# [
# { path: /nix/store/foo, references: [ /nix/store/foo, /nix/store/bar, /nix/store/baz ] },
# { path: /nix/store/bar, references: [ /nix/store/bar, /nix/store/baz ] },
# { path: /nix/store/baz, references: [ /nix/store/baz, /nix/store/tux ] },
# { path: /nix/store/tux, references: [ /nix/store/tux ] }
# ]
#
# To:
# {
# /nix/store/foo: [ /nix/store/bar, /nix/store/baz ],
# /nix/store/bar: [ /nix/store/baz ],
# /nix/store/baz: [ /nix/store/tux ] },
# /nix/store/tux: [ ]
# }
#
# Note that it drops self-references to avoid loops.
def make_lookup(closures):
lookup = {}
for closure in closures:
# paths often self-refer
nonreferential_paths = [ref for ref in closure['references'] if ref != closure['path']]
lookup[closure['path']] = nonreferential_paths
return lookup
class TestMakeLookup(unittest.TestCase):
def test_returns_lookp(self):
self.assertDictEqual(
make_lookup([
{
"path": "/nix/store/foo",
"references": [
"/nix/store/foo",
"/nix/store/bar"
]
},
{
"path": "/nix/store/bar",
"references": [
"/nix/store/bar",
"/nix/store/tux"
]
},
{
"path": "/nix/store/hello",
"references": [
]
}
]),
{
"/nix/store/foo": [ "/nix/store/bar" ],
"/nix/store/bar": [ "/nix/store/tux" ],
"/nix/store/hello": [ ],
}
)
# Convert:
#
# /nix/store/foo with
# {
# /nix/store/foo: [ /nix/store/bar, /nix/store/baz ],
# /nix/store/bar: [ /nix/store/baz ],
# /nix/store/baz: [ /nix/store/tux ] },
# /nix/store/tux: [ ]
# }
#
# To:
#
# {
# /nix/store/bar: {
# /nix/store/baz: {
# /nix/store/tux: {}
# }
# },
# /nix/store/baz: {
# /nix/store/tux: {}
# }
# }
subgraphs_cache = {}
def make_graph_segment_from_root(root, lookup):
global subgraphs_cache
children = {}
for ref in lookup[root]:
# make_graph_segment_from_root is a pure function, and will
# always return the same result based on a given input. Thus,
# cache computation.
#
# Python's assignment will use a pointer, preventing memory
# bloat for large graphs.
if ref not in subgraphs_cache:
debug("Subgraph Cache miss on {}".format(ref))
subgraphs_cache[ref] = make_graph_segment_from_root(ref, lookup)
else:
debug("Subgraph Cache hit on {}".format(ref))
children[ref] = subgraphs_cache[ref]
return children
class TestMakeGraphSegmentFromRoot(unittest.TestCase):
def test_returns_graph(self):
self.assertDictEqual(
make_graph_segment_from_root("/nix/store/foo", {
"/nix/store/foo": [ "/nix/store/bar" ],
"/nix/store/bar": [ "/nix/store/tux" ],
"/nix/store/tux": [ ],
"/nix/store/hello": [ ],
}),
{
"/nix/store/bar": {
"/nix/store/tux": {}
}
}
)
def test_returns_graph_tiny(self):
self.assertDictEqual(
make_graph_segment_from_root("/nix/store/tux", {
"/nix/store/foo": [ "/nix/store/bar" ],
"/nix/store/bar": [ "/nix/store/tux" ],
"/nix/store/tux": [ ],
}),
{}
)
# Convert a graph segment in to a popularity-counted dictionary:
#
# From:
# {
# /nix/store/foo: {
# /nix/store/bar: {
# /nix/store/baz: {
# /nix/store/tux: {}
# }
# }
# /nix/store/baz: {
# /nix/store/tux: {}
# }
# }
# }
#
# to:
# [
# /nix/store/foo: 1
# /nix/store/bar: 2
# /nix/store/baz: 4
# /nix/store/tux: 6
# ]
popularity_cache = {}
def graph_popularity_contest(full_graph):
global popularity_cache
popularity = defaultdict(int)
for path, subgraph in full_graph.items():
popularity[path] += 1
# graph_popularity_contest is a pure function, and will
# always return the same result based on a given input. Thus,
# cache computation.
#
# Python's assignment will use a pointer, preventing memory
# bloat for large graphs.
if path not in popularity_cache:
debug("Popularity Cache miss on {}", path)
popularity_cache[path] = graph_popularity_contest(subgraph)
else:
debug("Popularity Cache hit on {}", path)
subcontest = popularity_cache[path]
for subpath, subpopularity in subcontest.items():
debug("Calculating popularity for {}", subpath)
popularity[subpath] += subpopularity + 1
return popularity
class TestGraphPopularityContest(unittest.TestCase):
def test_counts_popularity(self):
self.assertDictEqual(
graph_popularity_contest({
"/nix/store/foo": {
"/nix/store/bar": {
"/nix/store/baz": {
"/nix/store/tux": {}
}
},
"/nix/store/baz": {
"/nix/store/tux": {}
}
}
}),
{
"/nix/store/foo": 1,
"/nix/store/bar": 2,
"/nix/store/baz": 4,
"/nix/store/tux": 6,
}
)
# Emit a list of packages by popularity, most first:
#
# From:
# [
# /nix/store/foo: 1
# /nix/store/bar: 1
# /nix/store/baz: 2
# /nix/store/tux: 2
# ]
#
# To:
# [ /nix/store/baz /nix/store/tux /nix/store/bar /nix/store/foo ]
def order_by_popularity(paths):
paths_by_popularity = defaultdict(list)
popularities = []
for path, popularity in paths.items():
popularities.append(popularity)
paths_by_popularity[popularity].append(path)
popularities = list(set(popularities))
popularities.sort()
flat_ordered = []
for popularity in popularities:
paths = paths_by_popularity[popularity]
paths.sort(key=package_name)
flat_ordered.extend(reversed(paths))
return list(reversed(flat_ordered))
class TestOrderByPopularity(unittest.TestCase):
def test_returns_in_order(self):
self.assertEqual(
order_by_popularity({
"/nix/store/foo": 1,
"/nix/store/bar": 1,
"/nix/store/baz": 2,
"/nix/store/tux": 2,
}),
[
"/nix/store/baz",
"/nix/store/tux",
"/nix/store/bar",
"/nix/store/foo"
]
)
def package_name(path):
parts = path.split('-')
start = parts.pop(0)
# don't throw away any data, so the order is always the same.
# even in cases where only the hash at the start has changed.
parts.append(start)
return '-'.join(parts)
def main():
filename = sys.argv[1]
key = sys.argv[2]
debug("Loading from {}", filename)
with open(filename) as f:
data = json.load(f)
# Data comes in as:
# [
# { path: /nix/store/foo, references: [ /nix/store/foo, /nix/store/bar, /nix/store/baz ] },
# { path: /nix/store/bar, references: [ /nix/store/bar, /nix/store/baz ] },
# { path: /nix/store/baz, references: [ /nix/store/baz, /nix/store/tux ] },
# { path: /nix/store/tux, references: [ /nix/store/tux ] }
# ]
#
# and we want to get out a list of paths ordered by how universally,
# important they are, ie: tux is referenced by every path, transitively
# so it should be #1
#
# [
# /nix/store/tux,
# /nix/store/baz,
# /nix/store/bar,
# /nix/store/foo,
# ]
graph = data[key]
debug("Finding roots from {}", key)
roots = find_roots(graph);
debug("Making lookup for {}", key)
lookup = make_lookup(graph)
full_graph = {}
for root in roots:
debug("Making full graph for {}", root)
full_graph[root] = make_graph_segment_from_root(root, lookup)
debug("Running contest")
contest = graph_popularity_contest(full_graph)
debug("Ordering by popularity")
ordered = order_by_popularity(contest)
debug("Checking for missing paths")
missing = []
for path in all_paths(graph):
if path not in ordered:
missing.append(path)
ordered.extend(missing)
print("\n".join(ordered))
if "--test" in sys.argv:
# Don't pass --test otherwise unittest gets mad
unittest.main(argv = [f for f in sys.argv if f != "--test" ])
else:
main()