From faf6b0ffb80db2bd8db9a535d3852b422082cde3 Mon Sep 17 00:00:00 2001
From: Neel Shah <neelshah.sa@gmail.com>
Date: Wed, 20 Dec 2023 16:44:38 +0100
Subject: [PATCH] Implement proper flushing logic on close for sessions and
 client reports

---
 CHANGELOG.md                              |  1 +
 sentry-ruby/lib/sentry-ruby.rb            | 11 ++++++-----
 sentry-ruby/lib/sentry/transport.rb       | 13 +++++++++++--
 sentry-ruby/spec/sentry/transport_spec.rb | 16 ++++++++++++++++
 sentry-ruby/spec/sentry_spec.rb           | 11 ++++++++++-
 5 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd1aa3633..8d0a788fd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@
 
     If your system serves heavy load, please let us know how this feature works for you!
 
+- Implement proper flushing logic on ``close`` for Client Reports and Sessions [#2206](https://github.com/getsentry/sentry-ruby/pull/2206)
 ## 5.15.2
 
 ### Bug Fixes
diff --git a/sentry-ruby/lib/sentry-ruby.rb b/sentry-ruby/lib/sentry-ruby.rb
index fd18b7756..900172f48 100644
--- a/sentry-ruby/lib/sentry-ruby.rb
+++ b/sentry-ruby/lib/sentry-ruby.rb
@@ -66,7 +66,7 @@ def exception_locals_tp
     end
 
     # @!attribute [rw] background_worker
-    #   @return [BackgroundWorker, nil]
+    #   @return [BackgroundWorker]
     attr_accessor :background_worker
 
     # @!attribute [r] session_flusher
@@ -233,10 +233,11 @@ def init(&block)
     #
     # @return [void]
     def close
-      if @background_worker
-        @background_worker.shutdown
-        @background_worker = nil
-      end
+      @background_worker.perform { get_current_client&.transport&.flush }
+      # session_flusher internally queues to the background worker too
+      @session_flusher&.flush
+
+      @background_worker.shutdown
 
       if @session_flusher
         @session_flusher.kill
diff --git a/sentry-ruby/lib/sentry/transport.rb b/sentry-ruby/lib/sentry/transport.rb
index 81cef5e97..2c3b1ca47 100644
--- a/sentry-ruby/lib/sentry/transport.rb
+++ b/sentry-ruby/lib/sentry/transport.rb
@@ -168,11 +168,20 @@ def record_lost_event(reason, item_type)
       @discarded_events[[reason, item_type]] += 1
     end
 
+    def flush
+      client_report_headers, client_report_payload = fetch_pending_client_report(force: true)
+      return unless client_report_headers
+
+      envelope = Envelope.new
+      envelope.add_item(client_report_headers, client_report_payload)
+      send_envelope(envelope)
+    end
+
     private
 
-    def fetch_pending_client_report
+    def fetch_pending_client_report(force: false)
       return nil unless @send_client_reports
-      return nil if @last_client_report_sent > Time.now - CLIENT_REPORT_INTERVAL
+      return nil if !force && @last_client_report_sent > Time.now - CLIENT_REPORT_INTERVAL
       return nil if @discarded_events.empty?
 
       discarded_events_hash = @discarded_events.map do |key, val|
diff --git a/sentry-ruby/spec/sentry/transport_spec.rb b/sentry-ruby/spec/sentry/transport_spec.rb
index fbad63039..d872d77b9 100644
--- a/sentry-ruby/spec/sentry/transport_spec.rb
+++ b/sentry-ruby/spec/sentry/transport_spec.rb
@@ -474,4 +474,20 @@
       end
     end
   end
+
+  describe "#flush" do
+    it "does not do anything without pending client reports" do
+      expect(subject).not_to receive(:send_envelope)
+      subject.flush
+    end
+
+    it "sends pending client reports" do
+      5.times { subject.record_lost_event(:ratelimit_backoff, 'error') }
+      3.times { subject.record_lost_event(:queue_overflow, 'transaction') }
+
+      expect(subject).to receive(:send_data)
+      subject.flush
+      expect(io.string).to match(/Sending envelope with items \[client_report\]/)
+    end
+  end
 end
diff --git a/sentry-ruby/spec/sentry_spec.rb b/sentry-ruby/spec/sentry_spec.rb
index 6f7cf78d8..a99bf1734 100644
--- a/sentry-ruby/spec/sentry_spec.rb
+++ b/sentry-ruby/spec/sentry_spec.rb
@@ -1005,7 +1005,6 @@
       it "calls background worker shutdown" do
         expect(described_class.background_worker).to receive(:shutdown)
         described_class.close
-        expect(described_class.background_worker).to eq(nil)
       end
 
       it "kills session flusher" do
@@ -1029,6 +1028,16 @@
         described_class.close
         expect(described_class.backpressure_monitor).to eq(nil)
       end
+
+      it "flushes transport" do
+        expect(described_class.get_current_client.transport).to receive(:flush)
+        described_class.close
+      end
+
+      it "flushes session flusher" do
+        expect(described_class.session_flusher).to receive(:flush)
+        described_class.close
+      end
     end
 
     it "can reinitialize closed SDK" do