Fixes an issue with storing collision values in page cache number array

Fixes #11888

community/kernel/src/main/java/org/neo4j/unsafe/impl/batchimport/cache/PageCachedNumberArrayFactory.java

@@ -39,7 +39,7 @@ public class PageCachedNumberArrayFactory extends NumberArrayFactory.Adapter
     private final PageCache pageCache;
     private final File storeDir;
 
-    PageCachedNumberArrayFactory( PageCache pageCache, File storeDir )
+    public PageCachedNumberArrayFactory( PageCache pageCache, File storeDir )
     {
         Objects.requireNonNull( pageCache );
         this.pageCache = pageCache;

community/kernel/src/main/java/org/neo4j/unsafe/impl/batchimport/cache/idmapping/string/StringCollisionValues.java

@@ -26,7 +26,7 @@
 
 import static java.lang.Integer.min;
 import static java.lang.Long.max;
-import static java.lang.Long.min;
+import static org.neo4j.io.pagecache.PageCache.PAGE_SIZE;
 
 /**
  * Stores {@link String strings} in a {@link ByteArray} provided by {@link NumberArrayFactory}. Each string can have different
@@ -41,7 +41,14 @@ public class StringCollisionValues implements CollisionValues
 
     public StringCollisionValues( NumberArrayFactory factory, long length )
     {
-        chunkSize = max( length, 10_000 );
+        // Let's have length (also chunk size) be divisible by PAGE_SIZE, such that our calculations below
+        // works for all NumberArray implementations.
+        if ( length % PAGE_SIZE != 0 )
+        {
+            length = ((length - 1) / PAGE_SIZE + 1) * PAGE_SIZE;
+        }
+
+        chunkSize = max( length, PAGE_SIZE );
         cache = factory.newDynamicByteArray( chunkSize, new byte[1] );
         current = cache.at( 0 );
     }
@@ -57,17 +64,10 @@ public long add( Object id )
             throw new IllegalArgumentException( string );
         }
 
-        long bytesLeftInThisChunk = bytesLeftInCurrentChunk();
-        if ( bytesLeftInThisChunk < Short.BYTES + 1 )
-        {
-            // There isn't enough space left in the current chunk to begin writing this value, move over to the next one
-            offset += chunkSize - (offset % chunkSize);
-            current = cache.at( offset );
-        }
-
         long startOffset = offset;
-        current.setShort( offset, 0, (short) length );
-        offset += Short.BYTES;
+        cache.setByte( offset++, 0, (byte) length );
+        cache.setByte( offset++, 0, (byte) (length >>> Byte.SIZE) );
+        current = cache.at( offset );
         for ( int i = 0; i < length; )
         {
             int bytesLeftToWrite = length - i;
@@ -87,18 +87,11 @@ public long add( Object id )
         return startOffset;
     }
 
-    private long bytesLeftInCurrentChunk()
-    {
-        long rest = offset % chunkSize;
-        return rest == 0 ? 0 : chunkSize - rest;
-    }
-
     @Override
     public Object get( long offset )
     {
+        int length = (cache.getByte( offset++, 0 ) & 0xFF) | ((cache.getByte( offset++, 0 ) & 0xFF) << Byte.SIZE);
         ByteArray array = cache.at( offset );
-        int length = array.getShort( offset, 0 ) & 0xFFFF;
-        offset += Short.BYTES;
         byte[] bytes = new byte[length];
         for ( int i = 0; i < length; )
         {

community/kernel/src/test/java/org/neo4j/unsafe/impl/batchimport/cache/idmapping/string/StringCollisionValuesTest.java

@@ -28,12 +28,17 @@
 
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.function.Function;
 
 import org.neo4j.test.Randoms;
+import org.neo4j.test.rule.PageCacheAndDependenciesRule;
 import org.neo4j.test.rule.RandomRule;
+import org.neo4j.test.rule.fs.DefaultFileSystemRule;
 import org.neo4j.unsafe.impl.batchimport.cache.NumberArrayFactory;
+import org.neo4j.unsafe.impl.batchimport.cache.PageCachedNumberArrayFactory;
 
 import static org.junit.Assert.assertEquals;
+import static org.neo4j.io.pagecache.PageCache.PAGE_SIZE;
 import static org.neo4j.unsafe.impl.batchimport.cache.NumberArrayFactory.AUTO_WITHOUT_PAGECACHE;
 import static org.neo4j.unsafe.impl.batchimport.cache.NumberArrayFactory.CHUNKED_FIXED_SIZE;
 import static org.neo4j.unsafe.impl.batchimport.cache.NumberArrayFactory.HEAP;
@@ -42,6 +47,8 @@
 @RunWith( Parameterized.class )
 public class StringCollisionValuesTest
 {
+    @Rule
+    public final PageCacheAndDependenciesRule storage = new PageCacheAndDependenciesRule( DefaultFileSystemRule::new, getClass() );
     @Rule
     public final RandomRule random = new RandomRule().withConfiguration( new Randoms.Default()
     {
@@ -53,19 +60,24 @@ public int stringMaxLength()
     } );
 
     @Parameters
-    public static Collection<NumberArrayFactory> data()
+    public static Collection<Function<PageCacheAndDependenciesRule,NumberArrayFactory>> data()
     {
-        return Arrays.asList( HEAP, OFF_HEAP, AUTO_WITHOUT_PAGECACHE, CHUNKED_FIXED_SIZE );
+        return Arrays.asList(
+                storage -> HEAP,
+                storage -> OFF_HEAP,
+                storage -> AUTO_WITHOUT_PAGECACHE,
+                storage -> CHUNKED_FIXED_SIZE,
+                storage -> new PageCachedNumberArrayFactory( storage.pageCache(), storage.directory().directory() ) );
     }
 
     @Parameter( 0 )
-    public NumberArrayFactory factory;
+    public Function<PageCacheAndDependenciesRule,NumberArrayFactory> factory;
 
     @Test
     public void shouldStoreAndLoadStrings()
     {
         // given
-        try ( StringCollisionValues values = new StringCollisionValues( factory, 10_000 ) )
+        try ( StringCollisionValues values = new StringCollisionValues( factory.apply( storage ), 10_000 ) )
         {
             // when
             long[] offsets = new long[100];
@@ -84,4 +96,27 @@ public void shouldStoreAndLoadStrings()
             }
         }
     }
+
+    @Test
+    public void shouldMoveOverToNextChunkOnNearEnd()
+    {
+        // given
+        try ( StringCollisionValues values = new StringCollisionValues( factory.apply( storage ), 10_000 ) )
+        {
+            char[] chars = new char[PAGE_SIZE - 3];
+            Arrays.fill( chars, 'a' );
+
+            // when
+            String string = String.valueOf( chars );
+            long offset = values.add( string );
+            String secondString = "abcdef";
+            long secondOffset = values.add( secondString );
+
+            // then
+            String readString = (String) values.get( offset );
+            assertEquals( string, readString );
+            String readSecondString = (String) values.get( secondOffset );
+            assertEquals( secondString, readSecondString );
+        }
+    }
 }