1 files changed, 79 insertions, 0 deletions
diff --git a/lib/gitlab/utils/deep_size.rb b/lib/gitlab/utils/deep_size.rb
new file mode 100644
index 00000000000..562cf09e249
--- /dev/null
+++ b/lib/gitlab/utils/deep_size.rb
@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+require 'objspace'
+
+module Gitlab
+  module Utils
+    class DeepSize
+      Error = Class.new(StandardError)
+      TooMuchDataError = Class.new(Error)
+
+      DEFAULT_MAX_SIZE = 1.megabyte
+      DEFAULT_MAX_DEPTH = 100
+
+      def initialize(root, max_size: DEFAULT_MAX_SIZE, max_depth: DEFAULT_MAX_DEPTH)
+        @root = root
+        @max_size = max_size
+        @max_depth = max_depth
+        @size = 0
+        @depth = 0
+
+        evaluate
+      end
+
+      def valid?
+        !too_big? && !too_deep?
+      end
+
+      private
+
+      def evaluate
+        add_object(@root)
+      rescue Error
+        # NOOP
+      end
+
+      def too_big?
+        @size > @max_size
+      end
+
+      def too_deep?
+        @depth > @max_depth
+      end
+
+      def add_object(object)
+        @size += ObjectSpace.memsize_of(object)
+        raise TooMuchDataError if @size > @max_size
+
+        add_array(object) if object.is_a?(Array)
+        add_hash(object) if object.is_a?(Hash)
+      end
+
+      def add_array(object)
+        with_nesting do
+          object.each do |n|
+            add_object(n)
+          end
+        end
+      end
+
+      def add_hash(object)
+        with_nesting do
+          object.each do |key, value|
+            add_object(key)
+            add_object(value)
+          end
+        end
+      end
+
+      def with_nesting
+        @depth += 1
+        raise TooMuchDataError if too_deep?
+
+        yield
+
+        @depth -= 1
+      end
+    end
+  end
+end