diff --git a/src/taskgraph/util/taskcluster.py b/src/taskgraph/util/taskcluster.py index b997f8a52..3b2245278 100644 --- a/src/taskgraph/util/taskcluster.py +++ b/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/{}") diff --git a/test/test_generator.py b/test/test_generator.py index 2dda5c671..be958680c 100644 --- a/test/test_generator.py +++ b/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