Add a method for updating a counter key using a buffer instead of a c…

…allback. Fix up asserts on prior test for setting counter key.
isabella232 · Apr 28, 2020 · 24b7bd1 · 24b7bd1
1 parent 721d172
commit 24b7bd1
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 1 deletion.
diff --git a/agrona/src/main/java/org/agrona/concurrent/status/CountersManager.java b/agrona/src/main/java/org/agrona/concurrent/status/CountersManager.java
@@ -400,6 +400,25 @@ public void setCounterKey(final int counterId, final Consumer<MutableDirectBuffe
         keyFunc.accept(new UnsafeBuffer(metaDataBuffer, metaDataOffset(counterId) + KEY_OFFSET, MAX_KEY_LENGTH));
     }
 
+    /**
+     * Set an {@link AtomicCounter} key based on counterId, copying the value from the supplied buffer.
+     *
+     * @param id        to be set
+     * @param keyBuffer containing the updated key
+     * @param offset    offset into buffer
+     * @param length    length of data to copy
+     */
+    public void setCounterKey(final int id, final DirectBuffer keyBuffer, final int offset, final int length)
+    {
+        if (length > MAX_KEY_LENGTH)
+        {
+            throw new IllegalArgumentException("Supplied key is too long: " + length + ", max: " + MAX_KEY_LENGTH);
+        }
+
+        metaDataBuffer.putBytes(metaDataOffset(id) + KEY_OFFSET, keyBuffer, offset, length);
+    }
+
+
     /**
      * Set an {@link AtomicCounter} label based on counterId.
      *

diff --git a/agrona/src/test/java/org/agrona/concurrent/status/CountersManagerTest.java b/agrona/src/test/java/org/agrona/concurrent/status/CountersManagerTest.java
@@ -287,13 +287,33 @@ public void shouldBeAbleToGetAndUpdateCounterLabel()
         assertThat(counter.label(), is("original label with update"));
     }
 
+    @Test
+    public void shouldBeAbleToGetAndUpdateCounterKeyUsingCallback()
+    {
+        final String originalKey = "original key";
+        final String updatedKey = "updated key";
+
+        final AtomicCounter counter = manager.newCounter(
+            "label", 101, (keyBuffer) -> keyBuffer.putStringUtf8(0, originalKey));
+
+        final StringKeyExtractor keyExtractor = new StringKeyExtractor(counter.id());
+
+        manager.forEach(keyExtractor);
+        assertThat(keyExtractor.key, is(originalKey));
+
+        manager.setCounterKey(counter.id(), (keyBuffer) -> keyBuffer.putStringUtf8(0, updatedKey));
+
+        manager.forEach(keyExtractor);
+        assertThat(keyExtractor.key, is(updatedKey));
+    }
 
     @Test
     public void shouldBeAbleToGetAndUpdateCounterKey()
     {
         final String originalKey = "original key";
         final String updatedKey = "updated key";
 
+
         final AtomicCounter counter = manager.newCounter(
             "label", 101, (keyBuffer) -> keyBuffer.putStringUtf8(0, originalKey));
 
@@ -303,8 +323,34 @@ public void shouldBeAbleToGetAndUpdateCounterKey()
 
         assertThat(keyExtractor.key, is(originalKey));
 
-        manager.setCounterKey(counter.id(), (keyBuffer) -> keyBuffer.putStringUtf8(0, updatedKey));
+        final UnsafeBuffer tempBuffer = new UnsafeBuffer(new byte[128]);
+        final int length = tempBuffer.putStringUtf8(0, updatedKey);
+
+        manager.setCounterKey(counter.id(), tempBuffer, 0, length);
+
+        manager.forEach(keyExtractor);
+        assertThat(keyExtractor.key, is(updatedKey));
+    }
+
+    @Test
+    public void shouldRejectOversizeKeys()
+    {
+        final String originalKey = "original key";
 
+        final AtomicCounter counter = manager.newCounter(
+            "label", 101, (keyBuffer) -> keyBuffer.putStringUtf8(0, originalKey));
+
+        final UnsafeBuffer tempBuffer = new UnsafeBuffer(new byte[256]);
+
+        try
+        {
+            manager.setCounterKey(counter.id(), tempBuffer, 0, MAX_KEY_LENGTH + 1);
+            fail("Should have thrown exception");
+        }
+        catch (IllegalArgumentException e)
+        {
+            assertTrue(true);
+        }
     }
 
     private static class StringKeyExtractor implements MetaData