DataOutputBuffer#scratchBuffer can use off-heap or on-heap memory as a means to control memory allocations

patch by David Capwell; reviewed by Caleb Rackliffe for CASSANDRA-16471

CHANGES.txt

@@ -1,4 +1,5 @@
 4.2
+ * DataOutputBuffer#scratchBuffer can use off-heap or on-heap memory as a means to control memory allocations (CASSANDRA-16471)
  * Add ability to read the TTLs and write times of the elements of a collection and/or UDT (CASSANDRA-8877)
  * Removed Python < 2.7 support from formatting.py (CASSANDRA-17694)
  * Cleanup pylint issues with pylexotron.py (CASSANDRA-17779)

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -289,10 +289,12 @@ public enum CassandraRelevantProperties
     /** property for the interval on which the repeated client warnings and diagnostic events about disk usage are ignored */
     DISK_USAGE_NOTIFY_INTERVAL_MS("cassandra.disk_usage.notify_interval_ms", Long.toString(TimeUnit.MINUTES.toMillis(30))),
 
+    /** Controls the type of bufffer (heap/direct) used for shared scratch buffers */
+    DATA_OUTPUT_BUFFER_ALLOCATE_TYPE("cassandra.dob.allocate_type"),
+
     // for specific tests
     ORG_APACHE_CASSANDRA_CONF_CASSANDRA_RELEVANT_PROPERTIES_TEST("org.apache.cassandra.conf.CassandraRelevantPropertiesTest"),
     ORG_APACHE_CASSANDRA_DB_VIRTUAL_SYSTEM_PROPERTIES_TABLE_TEST("org.apache.cassandra.db.virtual.SystemPropertiesTableTest"),
-
     ;
 
 
@@ -454,6 +456,40 @@ public enum CassandraRelevantProperties
         System.setProperty(key, Long.toString(value));
     }
 
+    /**
+     * Gets the value of a system property as a enum, calling {@link String#toUpperCase()} first.
+     *
+     * @param defaultValue to return when not defined
+     * @param <T> type
+     * @return enum value
+     */
+    public <T extends Enum<T>> T getEnum(T defaultValue) {
+        return getEnum(true, defaultValue);
+    }
+
+    /**
+     * Gets the value of a system property as a enum, optionally calling {@link String#toUpperCase()} first.
+     *
+     * @param toUppercase before converting to enum
+     * @param defaultValue to return when not defined
+     * @param <T> type
+     * @return enum value
+     */
+    public <T extends Enum<T>> T getEnum(boolean toUppercase, T defaultValue) {
+        String value = System.getProperty(key);
+        if (value == null)
+            return defaultValue;
+        return Enum.valueOf(defaultValue.getDeclaringClass(), toUppercase ? value.toUpperCase() : value);
+    }
+
+    /**
+     * Sets the value into system properties.
+     * @param value to set
+     */
+    public void setEnum(Enum<?> value) {
+        System.setProperty(key, value.name());
+    }
+
     public interface PropertyConverter<T>
     {
         T convert(String value);

src/java/org/apache/cassandra/db/commitlog/CommitLog.java

@@ -291,7 +291,7 @@ public class CommitLog implements CommitLogMBean
                 buffer.putInt((int) checksum.getValue());
 
                 // checksummed mutation
-                dos.write(dob.getData(), 0, size);
+                dos.write(dob.unsafeGetBufferAndFlip());
                 updateChecksum(checksum, buffer, buffer.position() - size, size);
                 buffer.putInt((int) checksum.getValue());
             }

src/java/org/apache/cassandra/db/rows/UnfilteredSerializer.java

@@ -196,7 +196,7 @@ public class UnfilteredSerializer
                 // We write the size of the previous unfiltered to make reverse queries more efficient (and simpler).
                 // This is currently not used however and using it is tbd.
                 out.writeUnsignedVInt(previousUnfilteredSize);
-                out.write(dob.getData(), 0, dob.getLength());
+                out.write(dob.unsafeGetBufferAndFlip());
             }
         }
         else

src/java/org/apache/cassandra/hints/HintsWriter.java

@@ -78,7 +78,7 @@ class HintsWriter implements AutoCloseable
         {
             // write the descriptor
             descriptor.serialize(dob);
-            ByteBuffer descriptorBytes = dob.buffer();
+            ByteBuffer descriptorBytes = dob.unsafeGetBufferAndFlip();
             updateChecksum(crc, descriptorBytes);
             channel.write(descriptorBytes);
 

src/java/org/apache/cassandra/io/util/BufferedDataOutputStreamPlus.java

@@ -66,6 +66,16 @@ public class BufferedDataOutputStreamPlus extends DataOutputStreamPlus
         this.buffer = buffer;
     }
 
+    protected BufferedDataOutputStreamPlus(int size)
+    {
+        this.buffer = allocate(size);
+    }
+
+    protected ByteBuffer allocate(int size)
+    {
+        return ByteBuffer.allocate(size);
+    }
+
     @Override
     public void write(byte[] b) throws IOException
     {

src/java/org/apache/cassandra/io/util/DataOutputBuffer.java

@@ -28,6 +28,8 @@ import com.google.common.base.Preconditions;
 import io.netty.util.concurrent.FastThreadLocal;
 import org.apache.cassandra.config.Config;
 
+import static org.apache.cassandra.config.CassandraRelevantProperties.DATA_OUTPUT_BUFFER_ALLOCATE_TYPE;
+
 /**
  * An implementation of the DataOutputStream interface using a FastByteArrayOutputStream and exposing
  * its buffer so copies can be avoided.
@@ -45,11 +47,13 @@ public class DataOutputBuffer extends BufferedDataOutputStreamPlus
      * Only recycle OutputBuffers up to 1Mb. Larger buffers will be trimmed back to this size.
      */
     private static final int MAX_RECYCLE_BUFFER_SIZE = Integer.getInteger(Config.PROPERTY_PREFIX + "dob_max_recycle_bytes", 1024 * 1024);
+    private enum AllocationType { DIRECT, ONHEAP }
+    private static final AllocationType ALLOCATION_TYPE = DATA_OUTPUT_BUFFER_ALLOCATE_TYPE.getEnum(AllocationType.DIRECT);
 
     private static final int DEFAULT_INITIAL_BUFFER_SIZE = 128;
 
     /**
-     * Scratch buffers used mostly for serializing in memory. It's important to call #recycle() when finished
+     * Scratch buffers used mostly for serializing in memory. It's important to call #close() when finished
      * to keep the memory overhead from being too large in the system.
      */
     public static final FastThreadLocal<DataOutputBuffer> scratchBuffer = new FastThreadLocal<DataOutputBuffer>()
@@ -59,29 +63,38 @@ public class DataOutputBuffer extends BufferedDataOutputStreamPlus
         {
             return new DataOutputBuffer()
             {
+                @Override
                 public void close()
                 {
-                    if (buffer.capacity() <= MAX_RECYCLE_BUFFER_SIZE)
+                    if (buffer != null && buffer.capacity() <= MAX_RECYCLE_BUFFER_SIZE)
                     {
                         buffer.clear();
                     }
                     else
                     {
-                        buffer = ByteBuffer.allocate(DEFAULT_INITIAL_BUFFER_SIZE);
+                        setBuffer(allocate(DEFAULT_INITIAL_BUFFER_SIZE));
                     }
                 }
+
+                @Override
+                protected ByteBuffer allocate(int size)
+                {
+                    return ALLOCATION_TYPE == AllocationType.DIRECT ?
+                           ByteBuffer.allocateDirect(size) :
+                           ByteBuffer.allocate(size);
+                }
             };
         }
     };
 
     public DataOutputBuffer()
     {
-        this(DEFAULT_INITIAL_BUFFER_SIZE);
+        super(DEFAULT_INITIAL_BUFFER_SIZE);
     }
 
     public DataOutputBuffer(int size)
     {
-        super(ByteBuffer.allocate(size));
+        super(size);
     }
 
     public DataOutputBuffer(ByteBuffer buffer)
@@ -158,9 +171,15 @@ public class DataOutputBuffer extends BufferedDataOutputStreamPlus
     {
         if (count <= 0)
             return;
-        ByteBuffer newBuffer = ByteBuffer.allocate(checkedArraySizeCast(calculateNewSize(count)));
+        ByteBuffer newBuffer = allocate(checkedArraySizeCast(calculateNewSize(count)));
         buffer.flip();
         newBuffer.put(buffer);
+        setBuffer(newBuffer);
+    }
+
+    protected void setBuffer(ByteBuffer newBuffer)
+    {
+        FileUtils.clean(buffer); // free if direct
         buffer = newBuffer;
     }
 
@@ -221,6 +240,18 @@ public class DataOutputBuffer extends BufferedDataOutputStreamPlus
         return result;
     }
 
+    /**
+     * Gets the underlying ByteBuffer and calls {@link ByteBuffer#flip()}.  This method is "unsafe" in the sense that
+     * it returns the underlying buffer, which may be modified by other methods after calling this method (or cleared on
+     * {@link #close()}). If the calling logic knows that no new calls to this object will happen after calling this
+     * method, then this method can avoid the copying done in {@link #asNewBuffer()}, and {@link #buffer()}.
+     */
+    public ByteBuffer unsafeGetBufferAndFlip()
+    {
+        buffer.flip();
+        return buffer;
+    }
+
     public byte[] getData()
     {
         assert buffer.arrayOffset() == 0;