Merge pull request #209 from mikeb01/set_counter_key

Add a method to update the key for a specific counter.

agrona/src/main/java/org/agrona/concurrent/status/CountersManager.java

@@ -389,6 +389,36 @@ public void setCounterLabel(final int counterId, final String label)
         putLabel(metaDataOffset(counterId), label);
     }
 
+    /**
+     * Set an {@link AtomicCounter} key based on counterId, using a consumer callback to update the key buffer.
+     *
+     * @param counterId to be set.
+     * @param keyFunc   callback used to set the key.
+     */
+    public void setCounterKey(final int counterId, final Consumer<MutableDirectBuffer> keyFunc)
+    {
+        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.
      *

agrona/src/test/java/org/agrona/concurrent/status/CountersManagerTest.java

@@ -287,6 +287,91 @@ 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));
+
+        final StringKeyExtractor keyExtractor = new StringKeyExtractor(counter.id());
+
+        manager.forEach(keyExtractor);
+
+        assertThat(keyExtractor.key, is(originalKey));
+
+        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 (final IllegalArgumentException e)
+        {
+            assertTrue(true);
+        }
+    }
+
+    private static final class StringKeyExtractor implements MetaData
+    {
+        private final int id;
+        private String key;
+
+        private StringKeyExtractor(final int id)
+        {
+            this.id = id;
+        }
+
+        public void accept(final int counterId, final int typeId, final DirectBuffer keyBuffer, final String label)
+        {
+            if (counterId == id)
+            {
+                key = keyBuffer.getStringUtf8(0);
+            }
+        }
+    }
+
     @Test
     public void shouldBeAbleToAppendLabel()
     {