Propagate OpenTelemetry context to async thread pool workers

.github/workflows/ci.yml

@@ -30,7 +30,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        ruby: ["3.0", "3.1", "3.2", "3.3"]
+        ruby: ["3.3", "3.4"]
 
     name: ${{ matrix.ruby }}
 

.github/workflows/release.yml

@@ -0,0 +1,54 @@
+name: Release
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Extract version
+        id: version
+        run: |
+          version=$(grep -oP '(?<=version = ")[^"]+' umbrellio-sequel-plugins.gemspec)
+          echo "tag=v$version" >> "$GITHUB_OUTPUT"
+
+      - name: Check if release already exists
+        id: check
+        run: |
+          if gh release view "${{ steps.version.outputs.tag }}" &>/dev/null; then
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+
+      - uses: ruby/setup-ruby@v1
+        if: steps.check.outputs.exists == 'false'
+        with:
+          ruby-version: "3.4"
+          bundler-cache: true
+
+      - name: Build gem
+        if: steps.check.outputs.exists == 'false'
+        run: gem build umbrellio-sequel-plugins.gemspec
+
+      - name: Push gem to RubyGems
+        if: steps.check.outputs.exists == 'false'
+        run: gem push *.gem
+        env:
+          GEM_HOST_API_KEY: ${{ secrets.GEM_HOST_API_KEY }}
+
+      - uses: softprops/action-gh-release@v3
+        if: steps.check.outputs.exists == 'false'
+        with:
+          tag_name: ${{ steps.version.outputs.tag }}
+          files: "*.gem"
+          generate_release_notes: true

.rubocop.yml

@@ -3,7 +3,7 @@ inherit_gem:
 
 AllCops:
   DisplayCopNames: true
-  TargetRubyVersion: 3.0
+  TargetRubyVersion: 3.3
 
 Naming/FileName:
   Exclude:

Gemfile

@@ -6,6 +6,7 @@ gemspec
 
 gem "async"
 gem "money"
+gem "opentelemetry-api"
 gem "pg"
 gem "pry"
 gem "rake"

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    umbrellio-sequel-plugins (0.18.0)
+    umbrellio-sequel-plugins (0.19.0)
       concurrent-ruby
       sequel
 
@@ -41,12 +41,15 @@ GEM
       concurrent-ruby (~> 1.0)
     json (2.7.2)
     language_server-protocol (3.17.0.3)
+    logger (1.7.0)
     method_source (1.1.0)
     minitest (5.25.1)
     money (6.19.0)
       i18n (>= 0.6.4, <= 2)
     mutex_m (0.2.0)
     nio4r (2.7.3)
+    opentelemetry-api (1.10.0)
+      logger
     parallel (1.26.3)
     parser (3.3.5.0)
       ast (~> 2.4.1)
@@ -140,6 +143,7 @@ PLATFORMS
 DEPENDENCIES
   async
   money
+  opentelemetry-api
   pg
   pry
   rake

lib/clickhouse/migrator.rb

@@ -18,13 +18,13 @@ def rollback(to: nil)
       migrator(target: target.to_i).run
     end
 
-    def migrator(**opts)
+    def migrator(**)
       Sequel::TimestampMigrator.new(
         DB,
         Rails.root.join("db/migrate/clickhouse"),
         table: :clickhouse_migrations,
         use_transactions: false,
-        **opts,
+        **,
       )
     end
   end

lib/sequel/extensions/concurrent_thread_pool.rb

@@ -11,8 +11,8 @@ def method_missing(...)
         __value.public_send(...)
       end
 
-      def respond_to_missing?(*args)
-        __value.respond_to?(*args)
+      def respond_to_missing?(*)
+        __value.respond_to?(*)
       end
 
       [:!, :==, :!=, :instance_eval, :instance_exec].each do |method|
@@ -24,10 +24,10 @@ def respond_to_missing?(*args)
 
     # Default proxy: schedules block via Concurrent::Future, blocks on first access.
     class Proxy < BaseProxy
-      def initialize(executor, &block)
+      def initialize(executor, &)
         super()
 
-        @future = Concurrent::Promises.future_on(executor, &block)
+        @future = Concurrent::Promises.future_on(executor, &)
       end
 
       def __value
@@ -119,8 +119,20 @@ def choose_executor(opts)
         end
       end
 
-      def async_run(&block)
-        async_job_class.new(async_thread_executor, &block)
+      def async_run(&)
+        if defined?(OpenTelemetry::Context)
+          otel_context = OpenTelemetry::Context.current
+
+          if otel_context.equal?(OpenTelemetry::Context::ROOT)
+            async_job_class.new(async_thread_executor, &)
+          else
+            async_job_class.new(async_thread_executor) do
+              OpenTelemetry::Context.with_current(otel_context, &)
+            end
+          end
+        else
+          async_job_class.new(async_thread_executor, &)
+        end
       end
     end
 

lib/sequel/extensions/migration_transaction_options.rb

@@ -11,9 +11,9 @@ module SimpleMigrationExtension
 end
 
 module MigratorExtension
-  def checked_transaction(migration, &block)
+  def checked_transaction(migration, &)
     if _use_transaction?(migration)
-      _transaction(migration, &block)
+      _transaction(migration, &)
     else
       yield
     end
@@ -38,11 +38,11 @@ def _use_transaction?(migration)
     end
   end
 
-  def _transaction(migration, &block)
+  def _transaction(migration, &)
     if migration.transaction_opts.nil?
-      db.transaction(&block)
+      db.transaction(&)
     else
-      db.transaction(migration.transaction_opts, &block)
+      db.transaction(migration.transaction_opts, &)
     end
   end
 end

lib/sequel/extensions/synchronize.rb

@@ -24,9 +24,9 @@ module Synchronize
     def synchronize_with(*args, timeout: 10, savepoint: false, skip_if_locked: false)
       key = lock_key_for(args)
 
-      transaction(savepoint: savepoint) do
+      transaction(savepoint:) do
         hash = key_hash(key)
-        if get_lock(key, hash, timeout: timeout, skip_if_locked: skip_if_locked)
+        if get_lock(key, hash, timeout:, skip_if_locked:)
           log_info("locked with #{key} (#{hash})")
           yield
         end

lib/sequel/plugins/upsert.rb

@@ -16,7 +16,7 @@ def upsert_dataset(target: primary_key)
       where_spec = cols.map { |x| Sequel::Plugins::Upsert.distinct_expr(table_name, x) }.reduce(:|)
 
       dataset.insert_conflict(
-        target: target,
+        target:,
         update: update_spec,
         update_where: where_spec,
       )
@@ -25,21 +25,21 @@ def upsert_dataset(target: primary_key)
     # Executes the upsert request
     #
     # @param row [Hash] values
-    # @param options [Hash] options
+    # @param target [Symbol] target column
     #
     # @example
-    #   User.upsert(name: "John", email: "jd@test.com", target: :email)
+    #   User.upsert({name: "John", email: "jd@test.com"}, target: :email)
     # @return [Sequel::Model]
-    def upsert(row, **options)
-      upsert_dataset(**options).insert(sequel_values(row))
+    def upsert(row, **)
+      upsert_dataset(**).insert(sequel_values(row))
     end
 
     # Executes the upsert request for multiple rows
     # @see #upsert
     # @see #upsert_dataset
-    def multi_upsert(rows, **options)
+    def multi_upsert(rows, **)
       rows = rows.map { |row| sequel_values(row) }
-      upsert_dataset(**options).multi_insert(rows)
+      upsert_dataset(**).multi_insert(rows)
     end
 
     # Returns formatted row values

lib/sequel/plugins/with_lock.rb

@@ -10,7 +10,7 @@ def with_lock(mode = "FOR NO KEY UPDATE", savepoint: true)
       @__locked = true
 
       begin
-        db.transaction(savepoint: savepoint) do
+        db.transaction(savepoint:) do
           lock!(mode)
           yield
         end

lib/tasks/sequel/archive_migrations.rake

@@ -15,7 +15,7 @@ namespace :sequel do
 
     migrations = Rails.root.glob(migrations_path).map do |file|
       filename = file.basename.to_s
-      { version: filename.to_i, filename: filename, source: file.read }
+      { version: filename.to_i, filename:, source: file.read }
     end
 
     conflict_options = {

lib/tasks/sequel/rollback_archived_migrations.rake

@@ -21,7 +21,7 @@ namespace :sequel do
 
       migrator_args = {
         table: args[:migration_table],
-        use_transactions: use_transactions,
+        use_transactions:,
         allow_missing_migration_files: false,
       }.compact
       migrator = Sequel::TimestampMigrator.new(DB, tmpdir, migrator_args)

lib/tasks/sequel/rollback_missing_migrations.rake

@@ -23,7 +23,7 @@ namespace :sequel do
     path = Rails.root.join("db/migrate")
     migrator_args = {
       table: args[:table],
-      use_transactions: use_transactions,
+      use_transactions:,
       allow_missing_migration_files: false,
     }.compact
     migrator = Sequel::TimestampMigrator.new(DB, path, migrator_args)

spec/extensions/concurrent_thread_pool_helpers.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require "concurrent"
+
+ASYNC_DB_URL = ENV.fetch("DB_URL", "postgres:///sequel_plugins")
+
+def make_concurrent_db(**opts)
+  Sequel.connect(ASYNC_DB_URL, **opts).tap { |d| d.extension(:concurrent_thread_pool) }
+end

spec/extensions/concurrent_thread_pool_opentelemetry_spec.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "opentelemetry-api"
+require_relative "concurrent_thread_pool_helpers"
+
+RSpec.describe "concurrent_thread_pool extension OpenTelemetry context propagation" do
+  let(:db) { make_concurrent_db }
+  let(:otel_key) { OpenTelemetry::Context.create_key("sequel-test") }
+
+  after { db.disconnect rescue nil }
+
+  it "does not propagate context when calling thread context is ROOT" do
+    captured = :not_set
+
+    db.send(:async_run) { captured = OpenTelemetry::Context.current.value(otel_key) }.__value
+
+    expect(captured).to be_nil
+  end
+
+  it "does not error when OpenTelemetry is defined without Context" do
+    db = make_concurrent_db
+    stub_const("OpenTelemetry", Object.new)
+
+    expect { db.send(:async_run) { :ok }.__value }.not_to raise_error
+  ensure
+    db.disconnect rescue nil
+  end
+
+  it "propagates calling-thread context to worker thread" do
+    captured = nil
+
+    OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+      db.send(:async_run) { captured = OpenTelemetry::Context.current.value(otel_key) }.__value
+    end
+
+    expect(captured).to eq("sentinel")
+  end
+
+  context "with single-thread executor" do
+    let(:db) { make_concurrent_db(num_async_threads: 1) }
+
+    it "restores context in worker thread after block completes" do
+      OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+        db.send(:async_run) { :done }.__value
+
+        value_after = Concurrent::Promises.future_on(db.async_thread_executor) do
+          OpenTelemetry::Context.current.value(otel_key)
+        end.value!
+
+        expect(value_after).to be_nil
+      end
+    end
+  end
+
+  it "propagates context even when block raises" do
+    captured = nil
+
+    OpenTelemetry::Context.with_value(otel_key, "sentinel") do
+      result = db.send(:async_run) do
+        captured = OpenTelemetry::Context.current.value(otel_key)
+        raise "boom"
+      end
+
+      expect { result.__value }.to raise_error(RuntimeError, "boom")
+      expect(captured).to eq("sentinel")
+    end
+  end
+end

spec/extensions/concurrent_thread_pool_spec.rb

@@ -1,12 +1,6 @@
 # frozen_string_literal: true
 
-require "concurrent"
-
-ASYNC_DB_URL = ENV.fetch("DB_URL", "postgres:///sequel_plugins")
-
-def make_concurrent_db(**opts)
-  Sequel.connect(ASYNC_DB_URL, **opts).tap { |d| d.extension(:concurrent_thread_pool) }
-end
+require_relative "concurrent_thread_pool_helpers"
 
 RSpec.describe "concurrent_thread_pool extension" do
   let(:db) { make_concurrent_db(**db_opts) }