Clean-up *types.Timestamp a bit

Needed for the future PR.
FerretDB · AlekSi · Sep 10, 2023 · Sep 4, 2023 · Sep 4, 2023 · Sep 4, 2023
commit 4f2add7afde6509402893c4ea00a2f2196b9c780
diff --git a/internal/handlers/sjson/schema_test.go b/internal/handlers/sjson/schema_test.go
@@ -98,7 +98,7 @@ func TestSchemaMarshalUnmarshal(t *testing.T) {
 						{"t": "object", "$s": {
 							"p": {
 								"arr": {"t": "array", "i": [
-									{"t": "int"}, 
+									{"t": "int"},
 									{"t": "timestamp"}
 								]},
 								"bar": {"t": "null"},
@@ -164,13 +164,13 @@ func TestSchemaMarshalUnmarshal(t *testing.T) {
 			json: `{
 				"p": {
 					"obj": {"t": "object", "$s": {
-						"p": {	
+						"p": {
 							"arr": {"t": "array", "i": [
 								{"t": "object", "$s": {}}, {"t": "object", "$s": {
 									"p": {"foo": {"t": "array", "i":[]}}, "$k": ["foo"]
 								}}
 							]},
-							"empty-arr": {"t": "array", "i":[]}				
+							"empty-arr": {"t": "array", "i":[]}
 						},
 						"$k": ["arr", "empty-arr"]
 					}}

diff --git a/internal/types/object_id_test.go b/internal/types/object_id_test.go
@@ -24,35 +24,35 @@ import (
 //nolint:paralleltest // we modify the global objectIDProcess
 func TestNewObjectID(t *testing.T) {
 	objectIDProcess = [5]byte{0x0b, 0xad, 0xc0, 0xff, 0xee}
-	ts := time.Date(2022, time.April, 13, 12, 44, 42, 0, time.UTC)
+	d := time.Date(2022, time.April, 13, 12, 44, 42, 0, time.UTC)
 
 	objectIDCounter.Store(0)
 	assert.Equal(
 		t,
 		ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0x00, 0x00, 0x01},
-		newObjectIDTime(ts),
+		newObjectIDTime(d),
 	)
 	assert.Equal(
 		t,
 		ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0x00, 0x00, 0x02},
-		newObjectIDTime(ts),
+		newObjectIDTime(d),
 	)
 
 	// test wraparound
 	objectIDCounter.Store(1<<24 - 2)
 	assert.Equal(
 		t,
 		ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0xff, 0xff, 0xff},
-		newObjectIDTime(ts),
+		newObjectIDTime(d),
 	)
 	assert.Equal(
 		t,
 		ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0x00, 0x00, 0x00},
-		newObjectIDTime(ts),
+		newObjectIDTime(d),
 	)
 	assert.Equal(
 		t,
 		ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0x00, 0x00, 0x01},
-		newObjectIDTime(ts),
+		newObjectIDTime(d),
 	)
 }
diff --git a/internal/types/timestamp.go b/internal/types/timestamp.go
@@ -19,30 +19,24 @@ import (
 	"time"
 )
 
-type (
-	// Timestamp represents BSON type Timestamp.
-	Timestamp int64
-)
-
-// timestampCounter is an ordinal number for timestamps in the system.
-var timestampCounter uint32
+// Timestamp represents BSON type Timestamp.
+type Timestamp uint64
 
-// NewTimestamp returns a timestamp from time and an increment.
-func NewTimestamp(t time.Time, c uint32) Timestamp {
-	sec := t.Unix()
-	sec <<= 32
-	sec |= int64(c)
-	return Timestamp(sec)
-}
+// timestampCounter is a process-wide timestamp counter.
+var timestampCounter atomic.Uint32
 
-// NextTimestamp returns a timestamp from time and an internal ops counter.
+// NextTimestamp returns the next timestamp for the given time value.
 func NextTimestamp(t time.Time) Timestamp {
-	c := atomic.AddUint32(&timestampCounter, 1)
-	return NewTimestamp(t, c)
+	// Technically, that should be a counter within a second, not a process-wide,
+	// but that's good enough for us.
+	c := uint64(timestampCounter.Add(1))
+
+	sec := uint64(t.Unix())
+	return Timestamp((sec << 32) | c)
 }
 
-// Time returns time.Time ignoring increment.
-func (t Timestamp) Time() time.Time {
-	t >>= 32
-	return time.Unix(int64(t), 0)
+// Time returns timestamp's time component.
+func (ts Timestamp) Time() time.Time {
+	sec := int64(ts >> 32)
+	return time.Unix(sec, 0).UTC()
 }
diff --git a/internal/types/timestamp_test.go b/internal/types/timestamp_test.go
@@ -0,0 +1,45 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package types
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+//nolint:paralleltest // we modify the global timestampCounter
+func TestNextTimestamp(t *testing.T) {
+	t.Run("UnixZero", func(t *testing.T) {
+		d := time.Unix(0, 0).UTC()
+
+		timestampCounter.Store(0)
+		assert.Equal(t, Timestamp(1), NextTimestamp(d))
+		assert.Equal(t, Timestamp(2), NextTimestamp(d))
+
+		assert.Equal(t, d, NextTimestamp(d).Time())
+	})
+
+	t.Run("Normal", func(t *testing.T) {
+		d := time.Date(2023, time.September, 12, 59, 44, 42, 0, time.UTC)
+
+		timestampCounter.Store(0)
+		assert.Equal(t, Timestamp(7278646209986691073), NextTimestamp(d))
+		assert.Equal(t, Timestamp(7278646209986691074), NextTimestamp(d))
+
+		assert.Equal(t, d, NextTimestamp(d).Time())
+	})
+}