Handle Hibernate gracefully.

Hibernate gets sequence values and caches by range
for assignment within hibernate, leading to potential
duplicates.

Make use of INCREMENT value of sequence to force
updating of ms time portion.

Author: mason-sharp

snowflake.c

@@ -172,7 +172,9 @@ snowflake_nextval(PG_FUNCTION_ARGS)
 	Buffer		buf;
 	Page		page;
 	HeapTupleData seqdatatuple;
+	HeapTuple	pgstuple;
 	Form_pg_sequence_data seq;
+	Form_pg_sequence	pgsform;
 	int64		result;
 	bool		logit = false;
 	Snowflake	flake;
@@ -243,9 +245,25 @@ snowflake_nextval(PG_FUNCTION_ARGS)
 		 * The clock either has not ticked or is behind. We need to make
 		 * sure that the flake doesn't move backwards and that we bump
 		 * it into the future should the count roll over.
+		 *
+		 * There is special handing for when ids are assigned in ranges
+		 * like for Hibernate. INCREMENT as part of the sequence definition
+		 * can be used. If the count portion of snowflake exceeds the
+		 * count mask + 1 (4096), it will cause the count to be set to 0
+		 * and the time ms portion is incremented.
 		 */
-		flake.sf_count++;
-		if ((flake.sf_count & SNOWFLAKE_COUNT_MASK) == 0)
+
+		pgstuple = SearchSysCache1(SEQRELID, ObjectIdGetDatum(relid));
+		if (HeapTupleIsValid(pgstuple))
+		{
+			pgsform = (Form_pg_sequence) GETSTRUCT(pgstuple);
+			elm->increment = pgsform->seqincrement;
+			ReleaseSysCache(pgstuple);
+		}
+		flake.sf_count += elm->increment;
+
+		/* Mask uses least signicant bits, so this is safe */
+		if (flake.sf_count > SNOWFLAKE_COUNT_MASK)
 		{
 			flake.sf_count = 0;
 			flake.sf_msec++;
@@ -533,6 +551,7 @@ init_sequence(Oid relid, SeqTable *p_elm, Relation *p_rel)
 		elm->lxid = InvalidLocalTransactionId;
 		elm->last_valid = false;
 		elm->last = elm->cached = 0;
+		elm->increment = 1;
 	}
 
 	/*

sql/conversion.sql

@@ -99,7 +99,57 @@ INSERT INTO t5 values (DEFAULT);
 SELECT x AS sf_val FROM t4 WHERE x > 42 \gset
 SELECT :sf_val < x AS is_growing FROM t5 WHERE x > 42;
 
+-- Test INCREMENT
+CREATE TABLE t6(x bigserial, y int);
+INSERT INTO t6 VALUES (DEFAULT, 10);
+INSERT INTO t6 VALUES (DEFAULT, 20);
+SELECT * FROM t6;
+INSERT INTO t6 VALUES (DEFAULT, 30);
+SELECT * FROM t6;
+SELECT snowflake.convert_sequence_to_snowflake('t6_x_seq');
+-- Get the count portion of the id.
+-- It happens within the same milisecond, should increment
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'));
+
+-- Should return 1 (the counter difference)
+SELECT ABS(snowflake.nextval('t6_x_seq') - snowflake.nextval('t6_x_seq'));
+
+-- If < 4096, will increment that amount
+ALTER SEQUENCE t6_x_seq INCREMENT 100 NO MAXVALUE;
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'));
+
+-- This will force the time ms portion to increment,
+-- the count portion should be 0
+ALTER SEQUENCE t6_x_seq INCREMENT 4096;
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'));
+
+-- should return 1 (the ms difference)
+SELECT ABS(snowflake.nextval('t6_x_seq') - snowflake.nextval('t6_x_seq')) >> 22;
+
+-- Test unreasonable value- still should get 0s
+ALTER SEQUENCE t6_x_seq INCREMENT 9999;
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'))
+UNION ALL
+SELECT snowflake.get_count(snowflake.nextval('t6_x_seq'));
+
+-- should return 1 (the ms difference)
+SELECT ABS(snowflake.nextval('t6_x_seq') - snowflake.nextval('t6_x_seq')) >> 22;
+
 -- Cleanup
-DROP TABLE t1,t2,t3,t4,t5 CASCADE;
+DROP TABLE t1,t2,t3,t4,t5,t6 CASCADE;
 DROP SEQUENCE favorite_seq;
 DROP EXTENSION snowflake;