2 files changed, 35 insertions, 4 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?
diff --git a/lib/tasks/db.rake b/lib/tasks/db.rake
index 6af6bb6fb..32039c31d 100644
--- a/lib/tasks/db.rake
+++ b/lib/tasks/db.rake
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require Rails.root.join('lib', 'mastodon', 'timestamp_ids')
+require_relative '../mastodon/snowflake'
 
 def each_schema_load_environment
   # If we're in development, also run this for the test environment.
@@ -63,13 +63,13 @@ namespace :db do
 
   task :define_timestamp_id do
     each_schema_load_environment do
-      Mastodon::TimestampIds.define_timestamp_id
+      Mastodon::Snowflake.define_timestamp_id
     end
   end
 
   task :ensure_id_sequences_exist do
     each_schema_load_environment do
-      Mastodon::TimestampIds.ensure_id_sequences_exist
+      Mastodon::Snowflake.ensure_id_sequences_exist
     end
   end
 end