Run sprite/numeric generation in threads

serveradmin/graphite/management/commands/cache_graphite.py

@@ -5,6 +5,7 @@
 
 import json
 import time
+from concurrent.futures import ThreadPoolExecutor
 from decimal import Decimal
 from io import BytesIO
 from os import mkdir
@@ -40,10 +41,15 @@ class Command(BaseCommand):
     def add_arguments(self, parser: CommandParser) -> None:
         parser.add_argument("--collections", nargs='*', type=str, help='Generate/update only these collections.')
         parser.add_argument("--query", type=str, help="Generate/update only objects matching this Serveradmin query.")
+        parser.add_argument("--threads", type=int, default=5, help="Generate n sprites/numerics concurrently.")
 
     def handle(self, *args, **options):
         """The entry point of the command"""
 
+        if options["threads"] < 1:
+            self.stderr.write(self.style.ERROR(f"--threads must be greater 0!"))
+            exit(1)
+
         start = time.time()
 
         sprite_params = settings.GRAPHITE_SPRITE_PARAMS
@@ -70,21 +76,24 @@ def handle(self, *args, **options):
             if options["query"]:
                 query_filter.update(**parse_query(options["query"]))
 
-            for server in Query(query_filter, ["hostname"]):
-                graph_table = collection.graph_table(server, sprite_params)
-                if graph_table:
-                    self.generate_sprite(collection_dir, graph_table, server)
-                    self.stdout.write(f"[{now()}] Generated sprite for {server['hostname']}")
-                self.cache_numerics(collection, server)
-                self.stdout.write(f"[{now()}] Updated numerics for {server['hostname']}")
+            futures = []
+            with ThreadPoolExecutor(options["threads"]) as executor:
+                for server in Query(query_filter, ["hostname"]):
+                    futures.append(executor.submit(self.generate_sprite, collection_dir, server, collection, sprite_params))
+                    futures.append(executor.submit(self.cache_numerics, collection, server))
 
             self.stdout.write(f"[{now()}] Finished collection {collection}")
 
         end = time.time()
         self.stdout.write(self.style.SUCCESS(f"Total time: {end - start:.2f} seconds."))
 
-    def generate_sprite(self, collection_dir, graph_table, server):
+    def generate_sprite(self, collection_dir, server, collection, sprite_params):
         """Generate sprites for the given server using the given collection"""
+
+        graph_table = collection.graph_table(server, sprite_params)
+        if not graph_table:
+            return
+
         graphs = [v2 for k1, v1 in graph_table for k2, v2 in v1]
         sprite_width = settings.GRAPHITE_SPRITE_WIDTH
         sprite_height = settings.GRAPHITE_SPRITE_HEIGHT
@@ -99,6 +108,8 @@ def generate_sprite(self, collection_dir, graph_table, server):
 
         sprite_img.save(collection_dir + '/' + server['hostname'] + '.png')
 
+        self.stdout.write(f"[{now()}] Generated sprite for {server['hostname']}")
+
     def cache_numerics(self, collection, server):
         """Generate sprites for the given server using the given collection"""
         for numeric in collection.numeric_set.all():
@@ -138,6 +149,8 @@ def cache_numerics(self, collection, server):
                     defaults={'value': Decimal(value)},
                 )
 
+        self.stdout.write(f"[{now()}] Updated numerics for {server['hostname']}")
+
     def get_from_graphite(self, params):
         """Make a GET request to Graphite with the given params"""
         password_mgr = HTTPPasswordMgrWithDefaultRealm()