1 files changed, 32 insertions, 1 deletions
diff --git a/lib/mastodon/timestamp_ids.rb b/lib/mastodon/snowflake.rb
index 3b048a50c..219e323d4 100644
--- a/lib/mastodon/timestamp_ids.rb
+++ b/lib/mastodon/snowflake.rb
@@ -1,8 +1,32 @@
 # frozen_string_literal: true
 
-module Mastodon::TimestampIds
+module Mastodon::Snowflake
   DEFAULT_REGEX = /timestamp_id\('(?<seq_prefix>\w+)'/
 
+  class Callbacks
+    def self.around_create(record)
+      now = Time.now.utc
+
+      if record.created_at.nil? || record.created_at >= now || record.created_at == record.updated_at
+        yield
+      else
+        record.id = Mastodon::Snowflake.id_at(record.created_at)
+        tries     = 0
+
+        begin
+          yield
+        rescue ActiveRecord::RecordNotUnique
+          raise if tries > 100
+
+          tries     += 1
+          record.id += rand(100)
+
+          retry
+        end
+      end
+    end
+  end
+
   class << self
     # Our ID will be composed of the following:
     # 6 bytes (48 bits) of millisecond-level timestamp
@@ -114,6 +138,13 @@ module Mastodon::TimestampIds
       end
     end
 
+    def id_at(timestamp)
+      id  = timestamp.to_i * 1000 + rand(1000)
+      id  = id << 16
+      id += rand(2**16)
+      id
+    end
+
     private
 
     def already_defined?