DRILL-998: Limit amount of memory used by drill C++ client API

2 files changed, 10 insertions, 4 deletions

diff --git a/contrib/native/client/src/include/drill/common.hpp b/contrib/native/client/src/include/drill/common.hpp
index 2113ce5e1..59734dc5b 100644
--- a/contrib/native/client/src/include/drill/common.hpp
+++ b/contrib/native/client/src/include/drill/common.hpp
@@ -33,6 +33,9 @@
 #define MAX_CONNECT_STR 4096
 #define MAX_SOCK_RD_BUFSIZE  1024
 
+#define MEM_CHUNK_SIZE 64*1024; // 64K
+#define MAX_MEM_ALLOC_SIZE 256*1024*1024; // 256 MB 
+
 #ifdef _DEBUG
 #define EXTRA_DEBUGGING
 #define CODER_DEBUGGING
@@ -48,6 +51,9 @@ typedef Byte_t * ByteBuf_t;
 class FieldMetadata;
 typedef boost::shared_ptr< std::vector<Drill::FieldMetadata*> > FieldDefPtr;
 
+class AllocatedBuffer;
+typedef AllocatedBuffer* AllocatedBufferPtr;
+
 typedef enum{
     QRY_SUCCESS=0,
     QRY_FAILURE=1,
diff --git a/contrib/native/client/src/include/drill/recordBatch.hpp b/contrib/native/client/src/include/drill/recordBatch.hpp
index e9298bf17..9a3df2b6f 100644
--- a/contrib/native/client/src/include/drill/recordBatch.hpp
+++ b/contrib/native/client/src/include/drill/recordBatch.hpp
@@ -836,10 +836,10 @@ class ValueVectorFactory{
 class DECLSPEC_DRILL_CLIENT RecordBatch{
     public:
 
-        //m_allocatedBuffer is the memory block allocated to hold the incoming RPC message. Record BAtches operate on
-        //slices of the allcoated buffer. The first slice (the first Field Batch), begins at m_buffer. Data in the
+        //m_allocatedBuffer is the memory block allocated to hold the incoming RPC message. Record Batches operate on
+        //slices of the allocated buffer. The first slice (the first Field Batch), begins at m_buffer. Data in the
         //allocated buffer before m_buffer is mostly the RPC header, and the QueryResult object.
-        RecordBatch(exec::shared::QueryResult* pResult, ByteBuf_t r, ByteBuf_t b)
+        RecordBatch(exec::shared::QueryResult* pResult, AllocatedBufferPtr r, ByteBuf_t b)
                 :m_fieldDefs(new(std::vector<Drill::FieldMetadata*>)){
             m_pQueryResult=pResult;
             m_pRecordBatchDef=&pResult->def();
@@ -892,7 +892,7 @@ class DECLSPEC_DRILL_CLIENT RecordBatch{
     private:
         const exec::shared::QueryResult* m_pQueryResult;
         const exec::shared::RecordBatchDef* m_pRecordBatchDef;
-        ByteBuf_t m_allocatedBuffer;
+        AllocatedBufferPtr m_allocatedBuffer;
         ByteBuf_t m_buffer;
         //build the current schema out of the field metadata
         FieldDefPtr m_fieldDefs;