Merge pull request rubyonjets#206 from tongueroo/global-memoist

cache aws service clients globally for lambda performance

lib/jets/aws_services.rb

@@ -10,59 +10,60 @@
 require "aws-sdk-sqs"
 
 module Jets::AwsServices
-  autoload :StackStatus, 'jets/aws_services/stack_status'
+  autoload :GlobalMemoist, 'jets/aws_services/global_memoist'
   autoload :S3Bucket, 'jets/aws_services/s3_bucket'
+  autoload :StackStatus, 'jets/aws_services/stack_status'
 
+  include GlobalMemoist
   include StackStatus
-  extend Memoist
 
   def apigateway
     Aws::APIGateway::Client.new
   end
-  memoize :apigateway
+  global_memoize :apigateway
 
   def cfn
     Aws::CloudFormation::Client.new
   end
-  memoize :cfn
+  global_memoize :cfn
 
   def dynamodb
     Aws::DynamoDB::Client.new
   end
-  memoize :dynamodb
+  global_memoize :dynamodb
 
   def lambda
     Aws::Lambda::Client.new
   end
-  memoize :lambda
+  global_memoize :lambda
 
   def logs
     Aws::CloudWatchLogs::Client.new
   end
-  memoize :logs
+  global_memoize :logs
 
   def s3
     Aws::S3::Client.new
   end
-  memoize :s3
+  global_memoize :s3
 
   def s3_resource
     Aws::S3::Resource.new
   end
-  memoize :s3_resource
+  global_memoize :s3_resource
 
   def sns
     Aws::SNS::Client.new
   end
-  memoize :sns
+  global_memoize :sns
 
   def sqs
     Aws::SQS::Client.new
   end
-  memoize :sqs
+  global_memoize :sqs
 
   def sts
     Aws::STS::Client.new
   end
-  memoize :sts
+  global_memoize :sts
 end

lib/jets/aws_services/global_memoist.rb

@@ -0,0 +1,57 @@
+# We cache the clients globally to avoid re-instantiating them again after the initial Lambda cold start.
+#
+# Based on: https://hashrocket.com/blog/posts/implementing-a-macro-in-ruby-for-memoization
+# Except we use a global variable for the cache. So we'll use the same client across
+# all instances as well as across Lambda executions after the cold-start. Example:
+#
+#   class Foo
+#     def s3
+#       Aws::S3::Client.new
+#     end
+#     global_memoize :s3
+#   end
+#
+#   foo1 = Foo.new
+#   foo2 = Foo.new
+#   foo1.s3
+#   foo2.s3 # same client as foo1
+#
+# A prewarmed request after a cold-start will still use the same s3 client instance since it uses the
+# global variable $__memo_methods as the cache.
+#
+module Jets::AwsServices
+  module GlobalMemoist
+    def self.included(klass)
+      klass.extend(Macros)
+    end
+
+    module Macros
+      def global_memoize(*methods)
+        const_defined?(:"GlobalMemoist#{self.object_id}") ?
+          const_get(:"GlobalMemoist#{self.object_id}") : const_set(:"GlobalMemoist#{self.object_id}", Module.new)
+        mod = const_get(:"GlobalMemoist#{self.object_id}")
+
+        mod.class_eval do
+          methods.each do |method|
+            define_method(method) do |skip_cache = false|
+              $__memo_methods ||= {}
+              if $__memo_methods.include?(method) && !skip_cache
+                $__memo_methods[method]
+              else
+                $__memo_methods[method] = super()
+              end
+            end
+          end
+        end
+        prepend mod
+      end
+
+      def global_memoize_class_method(*methods)
+        singleton_class.class_eval do
+          include GlobalMemoist
+          global_memoize(*methods)
+        end
+      end
+    end
+  end
+end