From 54306a974802c2e4d17e0980531e2d0dba08150a Mon Sep 17 00:00:00 2001
From: Adam Cooke <adam@krystal.io>
Date: Tue, 6 Feb 2024 11:53:53 +0000
Subject: [PATCH] fix: adds new connection pool which will discard failed
 clients

closes #2780
---
 config/initializers/postal.rb                 |  1 -
 lib/postal/message_db/connection_pool.rb      | 66 +++++++++++++++++++
 lib/postal/message_db/database.rb             | 10 ++-
 lib/postal/message_db/mysql.rb                | 39 -----------
 .../postal/message_db/connection_pool_spec.rb | 45 +++++++++++++
 5 files changed, 120 insertions(+), 41 deletions(-)
 create mode 100644 lib/postal/message_db/connection_pool.rb
 delete mode 100644 lib/postal/message_db/mysql.rb
 create mode 100644 spec/lib/postal/message_db/connection_pool_spec.rb

diff --git a/config/initializers/postal.rb b/config/initializers/postal.rb
index ca8e222..e5cfc9f 100644
--- a/config/initializers/postal.rb
+++ b/config/initializers/postal.rb
@@ -1,2 +1 @@
 require "postal"
-require "postal/message_db/mysql"
diff --git a/lib/postal/message_db/connection_pool.rb b/lib/postal/message_db/connection_pool.rb
new file mode 100644
index 0000000..579349c
--- /dev/null
+++ b/lib/postal/message_db/connection_pool.rb
@@ -0,0 +1,66 @@
+module Postal
+  module MessageDB
+    class ConnectionPool
+
+      attr_reader :connections
+
+      def initialize
+        @connections = []
+        @lock = Mutex.new
+      end
+
+      def use
+        connection = checkout
+        do_not_checkin = false
+        begin
+          yield connection
+        rescue Mysql2::Error => e
+          if e.message =~ /(lost connection|gone away|not connected)/i
+            # If the connection has failed for a connectivity reason
+            # we won't add it back in to the pool so that it'll reconnect
+            # next time.
+            do_not_checkin = true
+          end
+
+          raise
+        ensure
+          checkin(connection) unless do_not_checkin
+        end
+      end
+
+      private
+
+      def checkout
+        @lock.synchronize do
+          return @connections.pop unless @connections.empty?
+        end
+
+        add_new_connection
+        checkout
+      end
+
+      def checkin(connection)
+        @lock.synchronize do
+          @connections << connection
+        end
+      end
+
+      def add_new_connection
+        @lock.synchronize do
+          @connections << establish_connection
+        end
+      end
+
+      def establish_connection
+        Mysql2::Client.new(
+          host: Postal.config.message_db.host,
+          username: Postal.config.message_db.username,
+          password: Postal.config.message_db.password,
+          port: Postal.config.message_db.port,
+          encoding: Postal.config.message_db.encoding || "utf8mb4"
+        )
+      end
+
+    end
+  end
+end
diff --git a/lib/postal/message_db/database.rb b/lib/postal/message_db/database.rb
index e63ba7e..ad9f26e 100644
--- a/lib/postal/message_db/database.rb
+++ b/lib/postal/message_db/database.rb
@@ -2,6 +2,14 @@ module Postal
   module MessageDB
     class Database
 
+      class << self
+
+        def connection_pool
+          @connection_pool ||= ConnectionPool.new
+        end
+
+      end
+
       def initialize(organization_id, server_id)
         @organization_id = organization_id
         @server_id = server_id
@@ -339,7 +347,7 @@ module Postal
       end
 
       def with_mysql(&block)
-        MessageDB::MySQL.client(&block)
+        self.class.connection_pool.use(&block)
       end
 
       def build_where_string(attributes, joiner = ", ")
diff --git a/lib/postal/message_db/mysql.rb b/lib/postal/message_db/mysql.rb
deleted file mode 100644
index dd0f463..0000000
--- a/lib/postal/message_db/mysql.rb
+++ /dev/null
@@ -1,39 +0,0 @@
-module Postal
-  module MessageDB
-    module MySQL
-
-      # This exists here because it needs to be required when the application loads
-      # so that it isn't unloaded in development. If it was unloaded in development,
-      # it would be undesirable as we'd just end up with lots of connections.
-
-      def self.new_client
-        Mysql2::Client.new(
-          host: Postal.config.message_db.host,
-          username: Postal.config.message_db.username,
-          password: Postal.config.message_db.password,
-          port: Postal.config.message_db.port,
-          encoding: Postal.config.message_db.encoding || "utf8mb4"
-        )
-      end
-
-      @free_clients = []
-
-      def self.client(&block)
-        client = @free_clients.shift || new_client
-        return_value = nil
-        tries = 2
-        begin
-          return_value = block.call(client)
-        rescue Mysql2::Error => e
-          raise unless e.message =~ /(lost connection|gone away)/i && (tries -= 1) > 0
-
-          retry
-        ensure
-          @free_clients << client
-        end
-        return_value
-      end
-
-    end
-  end
-end
diff --git a/spec/lib/postal/message_db/connection_pool_spec.rb b/spec/lib/postal/message_db/connection_pool_spec.rb
new file mode 100644
index 0000000..73760e0
--- /dev/null
+++ b/spec/lib/postal/message_db/connection_pool_spec.rb
@@ -0,0 +1,45 @@
+require "rails_helper"
+
+describe Postal::MessageDB::ConnectionPool do
+
+  subject(:pool) { described_class.new }
+
+  describe '#use' do
+    it 'yields a connection' do
+      counter = 0
+      pool.use do |connection|
+        expect(connection).to be_a Mysql2::Client
+        counter += 1
+      end
+      expect(counter).to eq 1
+    end
+
+    it 'checks in a connection after the block has executed' do
+      connection = nil
+      pool.use do |c|
+        expect(pool.connections).to be_empty
+        connection = c
+      end
+      expect(pool.connections).to eq [connection]
+    end
+
+    it 'checks in a connection if theres an error in the block' do
+      expect do
+        pool.use do |c|
+          raise StandardError
+        end
+      end.to raise_error StandardError
+      expect(pool.connections).to match [kind_of(Mysql2::Client)]
+    end
+
+    it 'does not check in connections when there is a connection error' do
+      expect do
+        pool.use do
+          raise Mysql2::Error, "lost connection to server"
+        end
+      end.to raise_error Mysql2::Error
+      expect(pool.connections).to eq []
+    end
+  end
+
+end