RUBY-245 Unique object id with timestamp

ext/cbson/cbson.c

@@ -881,7 +881,7 @@ static VALUE method_deserialize(VALUE self, VALUE bson) {
     return elements_to_hash(buffer, remaining);
 }
 
-static VALUE objectid_generate(VALUE self)
+static VALUE objectid_generate(int argc, VALUE* args, VALUE self)
 {
     VALUE oid;
     unsigned char oid_bytes[12];
@@ -889,7 +889,11 @@ static VALUE objectid_generate(VALUE self)
     unsigned short pid;
     int i;
 
-    t = htonl((int)time(NULL));
+    if(argc == 0 || (argc == 1 && *args == Qnil)) {
+        t = htonl((int)time(NULL));
+    } else {
+        t = htonl(FIX2INT(rb_funcall(*args, rb_intern("to_i"), 0)));
+    }
     MEMCPY(&oid_bytes, &t, unsigned char, 4);
 
     MEMCPY(&oid_bytes[4], hostname_digest, unsigned char, 3);
@@ -969,7 +973,7 @@ void Init_cbson() {
     Digest = rb_const_get(rb_cObject, rb_intern("Digest"));
     DigestMD5 = rb_const_get(Digest, rb_intern("MD5"));
 
-    rb_define_method(ObjectId, "generate", objectid_generate, 0);
+    rb_define_method(ObjectId, "generate", objectid_generate, -1);
 
     if (gethostname(hostname, MAX_HOSTNAME_LENGTH) != 0) {
         rb_raise(rb_eRuntimeError, "failed to get hostname");

lib/bson/types/object_id.rb

@@ -33,6 +33,8 @@ module BSON
     @@lock  = Mutex.new
     @@index = 0
 
+    attr_accessor :data
+
     # Create a new object id. If no parameter is given, an id corresponding
     # to the ObjectId BSON data type will be created. This is a 12-byte value
     # consisting of a 4-byte timestamp, a 3-byte machine id, a 2-byte process id,
@@ -40,12 +42,15 @@ module BSON
     #
     # @param [Array] data should be an array of bytes. If you want
     #   to generate a standard MongoDB object id, leave this argument blank.
-    def initialize(data=nil)
-      @data = data || generate
+    #
+    # @option opts :data (nil) An array of bytes to use as the object id.
+    # @option opts :time (nil) The value of this object ids timestamp. Note that
+    #   the remaining bytes will consist of the standard machine id, pid, and counter. If
+    #   you need a zeroed timestamp, used ObjectId.from_time.
+    def initialize(data=nil, time=nil)
+      @data = data || generate(time)
     end
 
-    attr_accessor :data
-
     # Determine if the supplied string is legal. Legal strings will
     # consist of 24 hexadecimal characters.
     #
@@ -62,14 +67,23 @@ module BSON
     #
     # @param [Time] time a utc time to encode as an object id.
     #
+    # @option opts [:unique] (false) If false, the object id's bytes
+    #   succeeding the timestamp will be zeroed; if true, they'll
+    #   consist of the standard machine id, pid, and counter.
+    #
     # @return [Mongo::ObjectId]
     #
     # @example Return all document created before Jan 1, 2010.
     #   time = Time.utc(2010, 1, 1)
     #   time_id = ObjectId.from_time(time)
     #   collection.find({'_id' => {'$lt' => time_id}})
-    def self.from_time(time)
-      self.new([time.to_i,0,0].pack("NNN").unpack("C12"))
+    def self.from_time(time, opts={})
+      unique = opts.fetch(:unique, false)
+      if unique
+        self.new(nil, time)
+      else
+        self.new([time.to_i,0,0].pack("NNN").unpack("C12"))
+      end
     end
 
     # Adds a primary key to the given document if needed.
@@ -166,12 +180,16 @@ module BSON
     private
 
     # This gets overwritten by the C extension if it loads.
-    def generate
+    def generate(oid_time=nil)
       oid = ''
 
       # 4 bytes current time
-      time = Time.new.to_i
-      oid += [time].pack("N")
+      if oid_time
+        t = oid_time.to_i
+      else
+        t = Time.new.to_i
+      end
+      oid += [t].pack("N")
 
       # 3 bytes machine
       oid += Digest::MD5.digest(Socket.gethostname)[0, 3]

test/bson/object_id_test.rb

@@ -126,9 +126,22 @@ class ObjectIdTest < Test::Unit::TestCase
     time = Time.now.utc
     id   = ObjectId.from_time(time)
 
+    assert id.to_a[4, 8].all? {|byte| byte == 0 }
     assert_equal time.to_i, id.generation_time.to_i
   end
 
+  def test_from_time_unique
+    time = Time.now.utc
+    id   = ObjectId.from_time(time, :unique => true)
+
+    mac_id = Digest::MD5.digest(Socket.gethostname)[0, 3].unpack("C3")
+    assert_equal id.to_a[4, 3], mac_id
+    assert_equal time.to_i, id.generation_time.to_i
+
+    id2 = ObjectId.new(nil, time)
+    assert_equal time.to_i, id2.generation_time.to_i
+  end
+
   def test_json
     id = ObjectId.new
     assert_equal "{\"$oid\": \"#{id}\"}", id.to_json