From d2b61976a4f8f73286906e2a6884d836a11fe4fb Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Mon, 8 Jul 2024 13:03:10 -0700
Subject: [PATCH] Share data loader to across asyncio boto sessions (#40658)

By default, a botocore session creates and caches an instance of JSONDecoder which
consumes a lot of memory.  This issue was reported here https://github.com/boto/botocore/issues/3078.
In the context of triggers which use boto sessions, this can result in excessive
memory usage and as a result reduced capacity on the triggerer.  We can reduce
memory footprint by sharing the loader instance across the sessions.
---
 airflow/providers/amazon/aws/hooks/base_aws.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py
index 0d88ded9b9fea..0d07bb16494f2 100644
--- a/airflow/providers/amazon/aws/hooks/base_aws.py
+++ b/airflow/providers/amazon/aws/hooks/base_aws.py
@@ -70,6 +70,19 @@
 
     from airflow.models.connection import Connection  # Avoid circular imports.
 
+_loader = botocore.loaders.Loader()
+"""
+botocore data loader to be used with async sessions
+
+By default, a botocore session creates and caches an instance of JSONDecoder which
+consumes a lot of memory.  This issue was reported here https://github.com/boto/botocore/issues/3078.
+In the context of triggers which use boto sessions, this can result in excessive
+memory usage and as a result reduced capacity on the triggerer.  We can reduce
+memory footprint by sharing the loader instance across the sessions.
+
+:meta private:
+"""
+
 
 class BaseSessionFactory(LoggingMixin):
     """
@@ -155,7 +168,9 @@ def _apply_session_kwargs(self, session):
     def get_async_session(self):
         from aiobotocore.session import get_session as async_get_session
 
-        return async_get_session()
+        session = async_get_session()
+        session.register_component("data_loader", _loader)
+        return session
 
     def create_session(
         self, deferrable: bool = False