fix: don't reuse the parent's HTTP connections in forked children (bug 2088960)

src/taskgraph/util/taskcluster.py

@@ -130,6 +130,16 @@ def get_session():
     return requests_retry_session(retries=5)
 
 
+def _reset_sessions_in_child():
+    # Connections pooled before a fork would be shared with the parent.
+    get_session.cache_clear()
+    get_taskcluster_client.cache_clear()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_sessions_in_child)
+
+
 def get_artifact_url(task_id, path, use_proxy=False):
     url = get_root_url(block_proxy=not use_proxy)
     artifact_tmpl = liburls.api(url, "queue", "v1", "task/{}/artifacts/{}")

test/test_generator.py

@@ -17,6 +17,7 @@
 from taskgraph import generator, graph
 from taskgraph.generator import Kind, load_tasks_for_kind, load_tasks_for_kinds
 from taskgraph.loader.default import loader as default_loader
+from taskgraph.util import taskcluster as tc_util
 from taskgraph.util.schema import SchemaValidationError
 
 linuxonly = pytest.mark.skipif(
@@ -446,6 +447,39 @@ def record_gc(self, *args, **kwargs):
     assert record.read_text() == "False"
 
 
+@linuxonly
+def test_forked_children_get_fresh_sessions(mocker, tmp_path, maketgg):
+    "Children loading kinds don't share the parent's HTTP sessions."
+    record = tmp_path / "sessions"
+    load_tasks = Kind.load_tasks
+    session = tc_util.get_session()
+    queue = tc_util.get_taskcluster_client("queue")
+
+    def record_sessions(self, *args, **kwargs):
+        record.write_text(
+            repr(
+                (
+                    tc_util.get_session.cache_info().currsize,
+                    tc_util.get_taskcluster_client.cache_info().currsize,
+                    tc_util.get_session() is session,
+                    tc_util.get_taskcluster_client("queue") is queue,
+                )
+            )
+        )
+        return load_tasks(self, *args, **kwargs)
+
+    mocker.patch.object(Kind, "load_tasks", record_sessions)
+    tgg = maketgg()
+    tgg._load_tasks_forked(
+        {k.name: k for k in tgg._load_kinds(tgg.graph_config)},
+        tgg.kind_graph,
+        tgg.parameters,
+    )
+    assert record.read_text() == repr((0, 0, False, False))
+    assert tc_util.get_session() is session
+    assert tc_util.get_taskcluster_client("queue") is queue
+
+
 def test_load_tasks_for_kind(monkeypatch):
     """
     `load_tasks_for_kinds` will load the tasks for the provided kind