Acquire semian resource for each retry in Active Record adapters

.devcontainer/Dockerfile

@@ -1,5 +1,5 @@
 ARG RUBY_VERSION=3.4.3
-FROM ruby:${RUBY_VERSION} as base
+FROM ruby:${RUBY_VERSION} AS base
 
 # Avoid warnings by switching to noninteractive
 ENV DEBIAN_FRONTEND=noninteractive
@@ -22,7 +22,7 @@ RUN apt-get -o Acquire::Max-FutureTime=86400 update \
 
 WORKDIR /opt
 RUN git clone --depth 1 https://github.com/brendangregg/FlameGraph
-ENV PATH /opt/FlameGraph:$PATH
+ENV PATH="/opt/FlameGraph:$PATH"
 
 # Switch back to dialog for any ad-hoc use of apt-get
 ENV DEBIAN_FRONTEND=dialog
@@ -31,4 +31,7 @@ COPY Gemfile* semian.gemspec /workspace/
 COPY lib /workspace/lib
 
 WORKDIR /workspace
+
+# Limit grpc native extension build parallelism to avoid running out of memory
+ENV GRPC_RUBY_BUILD_PROCS=4
 RUN bundle install

.devcontainer/devcontainer.json

@@ -7,30 +7,29 @@
     "shutdownAction": "stopCompose",
     "postCreateCommand": "bundle install",
 
-    "settings": {
-        "C_Cpp.updateChannel": "Insiders",
-        "ruby.format": "rubocop",
-        "ruby.codeCompletion": "rcodetools",
-        "ruby.useLanguageServer": true,
-        "ruby.useBundler": true,
-        "ruby.intellisense": "rubyLocate",
-        "ruby.lint": {
-            "rubocop": true,
-        },
-        "ruby.lintDebounceTime": 500,
-        "[ruby]": {
-            "editor.insertSpaces": true,
-			"editor.tabSize": 2,
-            "editor.formatOnSaveTimeout": 1500,
-            //TODO:paranoidaditya remove comment after repo is formatted correctly
-            // "editor.formatOnSave": true,
-            "editor.defaultFormatter": "rebornix.ruby",
-        },
-        "terminal.integrated.shell.linux": "/bin/bash"
-    },
-
-    "extensions": [
-        "ms-vscode.cpptools",
-        "rebornix.ruby"
-    ]
+    "customizations": {
+        "vscode": {
+            "settings": {
+                "C_Cpp.updateChannel": "Insiders",
+                "ruby.format": "rubocop",
+                "ruby.codeCompletion": "rcodetools",
+                "ruby.useLanguageServer": true,
+                "ruby.useBundler": true,
+                "ruby.intellisense": "rubyLocate",
+                "ruby.lint": {
+                    "rubocop": true
+                },
+                "ruby.lintDebounceTime": 500,
+                "[ruby]": {
+                    "editor.insertSpaces": true,
+                    "editor.tabSize": 2,
+                    "editor.formatOnSaveTimeout": 1500
+                },
+                "terminal.integrated.shell.linux": "/bin/bash"
+            },
+            "extensions": [
+                "ms-vscode.cpptools"
+            ]
+        }
+    }
 }

.devcontainer/docker-compose.yml

@@ -1,6 +1,5 @@
 ---
 
-version: "3.7"
 services:
   semian: &base
     container_name: semian
@@ -55,25 +54,21 @@ services:
 
   toxiproxy:
     image: ghcr.io/shopify/toxiproxy:2.12.0
-    container_name: toxiproxy
     depends_on:
       - redis
       - mysql
 
   redis:
-    container_name: redis
     image: redis:latest
     command: redis-server
 
   mysql:
-    container_name: mysql
     image: mysql:9.3
     environment:
       MYSQL_ROOT_PASSWORD: root
       MYSQL_ROOT_HOST: "%"
 
   postgres:
-    container_name: postgres
     image: postgres:15
     environment:
       POSTGRES_PASSWORD: root

lib/semian/activerecord_adapter.rb

@@ -54,26 +54,6 @@ def initialize(*options)
       super
     end
 
-    if ActiveRecord.version >= Gem::Version.new("8.2.a")
-      def execute_intent(intent)
-        return super if self.class.query_allowlisted?(intent.processed_sql)
-
-        acquire_semian_resource(adapter: semian_adapter_name, scope: :query) do
-          super
-        end
-      end
-    else
-      def raw_execute(sql, *args, **kwargs, &block)
-        if self.class.query_allowlisted?(sql)
-          super
-        else
-          acquire_semian_resource(adapter: semian_adapter_name, scope: :query) do
-            super(sql, *args, **kwargs, &block)
-          end
-        end
-      end
-    end
-
     def active?
       acquire_semian_resource(adapter: semian_adapter_name, scope: :ping) do
         super
@@ -110,6 +90,28 @@ def connect(*args)
       end
     end
 
+    if ActiveRecord.version >= Gem::Version.new("8.2.a")
+      def perform_query(raw_connection, intent)
+        return super if self.class.query_allowlisted?(intent.processed_sql)
+
+        acquire_semian_resource(adapter: semian_adapter_name, scope: :query) do
+          super
+        rescue => e
+          raise translate_exception_class(e, intent.processed_sql, nil)
+        end
+      end
+    else
+      def raw_execute(sql, *args, **kwargs, &block)
+        if self.class.query_allowlisted?(sql)
+          super
+        else
+          acquire_semian_resource(adapter: semian_adapter_name, scope: :query) do
+            super(sql, *args, **kwargs, &block)
+          end
+        end
+      end
+    end
+
     def semian_adapter_name
       raise NotImplementedError, "#{self.class} must implement an `semian_adapter_name` method"
     end

test/adapters/activerecord_adapter_shared_tests.rb

@@ -347,6 +347,31 @@ def test_circuit_open_errors_do_not_trigger_the_circuit_breaker
     end
   end
 
+  def test_semian_counts_each_retry_attempt_individually
+    @adapter = new_adapter(
+      connection_retries: 1,
+      semian: SEMIAN_OPTIONS.merge(error_threshold: 3),
+    )
+    @adapter.connect!
+
+    query_acquisition_count = 0
+    subscriber = Semian.subscribe do |event, _resource, scope, _adapter|
+      if event == :success && scope == :query
+        query_acquisition_count += 1
+      end
+    end
+
+    # Close the raw connection to trigger a retryable connection error
+    @adapter.send(:raw_connection).close
+
+    value = @adapter.query_value("SELECT 1;")
+
+    assert_equal(2, value)
+    assert_equal(2, query_acquisition_count)
+  ensure
+    Semian.unsubscribe(subscriber)
+  end
+
   private
 
   def adapter_class