ICE-6852 - allow OutputStream to marshal to user-supplied buffer

6 files changed, 94 insertions, 28 deletions

diff --git a/cpp/include/Ice/Buffer.h b/cpp/include/Ice/Buffer.h
index 5274950ede7..3cde5701a99 100644
--- a/cpp/include/Ice/Buffer.h
+++ b/cpp/include/Ice/Buffer.h
@@ -86,8 +86,6 @@ public:
 
         void resize(size_type n) // Inlined for performance reasons.
         {
-            assert(!_buf || _capacity > 0);
-
             if(n == 0)
             {
                 clear();
@@ -101,14 +99,12 @@ public:
 
         void reset()
         {
-            assert(!_buf || _capacity > 0);
-
             if(_size > 0 && _size * 2 < _capacity)
             {
                 //
                 // If the current buffer size is smaller than the
                 // buffer capacity, we shrink the buffer memory to the
-                // current size. This is to avoid holding on too much
+                // current size. This is to avoid holding onto too much
                 // memory if it's not needed anymore.
                 //
                 if(++_shrinkCounter > 2)
@@ -152,6 +148,7 @@ public:
         size_type _size;
         size_type _capacity;
         int _shrinkCounter;
+        bool _owned;
     };
 
     Container b;
diff --git a/cpp/include/Ice/OutputStream.h b/cpp/include/Ice/OutputStream.h
index ed2ddc3783c..ee45b8b4429 100644
--- a/cpp/include/Ice/OutputStream.h
+++ b/cpp/include/Ice/OutputStream.h
@@ -49,6 +49,13 @@ public:
     //
     OutputStream(const CommunicatorPtr&, const EncodingVersion&);
 
+    //
+    // This constructor uses the given communicator and encoding version. The byte pair denotes
+    // application-supplied memory that the stream uses as its initial marshaling buffer. The
+    // stream will reallocate if the size of the marshaled data exceeds the application's buffer.
+    //
+    OutputStream(const CommunicatorPtr&, const EncodingVersion&, const std::pair<const Byte*, const Byte*>&);
+
     ~OutputStream()
     {
         // Inlined for performance reasons.
diff --git a/cpp/src/Ice/Buffer.cpp b/cpp/src/Ice/Buffer.cpp
index fe6c5490f78..33cbaeabc34 100644
--- a/cpp/src/Ice/Buffer.cpp
+++ b/cpp/src/Ice/Buffer.cpp
@@ -25,31 +25,36 @@ IceInternal::Buffer::Container::Container() :
     _buf(0),
     _size(0),
     _capacity(0),
-    _shrinkCounter(0)
+    _shrinkCounter(0),
+    _owned(true)
 {
 }
 
 IceInternal::Buffer::Container::Container(const_iterator beg, const_iterator end) :
     _buf(const_cast<iterator>(beg)),
     _size(end - beg),
-    _capacity(0),
-    _shrinkCounter(0)
+    _capacity(end - beg),
+    _shrinkCounter(0),
+    _owned(false)
 {
 }
 
 IceInternal::Buffer::Container::Container(const vector<value_type>& v) :
-    _capacity(0),
     _shrinkCounter(0)
 {
     if(v.empty())
     {
         _buf = 0;
         _size = 0;
+        _capacity = 0;
+        _owned = true;
     }
     else
     {
         _buf = const_cast<value_type*>(&v[0]);
         _size = v.size();
+        _capacity = _size;
+        _owned = false;
     }
 }
 
@@ -61,24 +66,27 @@ IceInternal::Buffer::Container::Container(Container& other, bool adopt)
         _size = other._size;
         _capacity = other._capacity;
         _shrinkCounter = other._shrinkCounter;
+        _owned = other._owned;
 
         other._buf = 0;
         other._size = 0;
         other._capacity = 0;
         other._shrinkCounter = 0;
+        other._owned = true;
     }
     else
     {
         _buf = other._buf;
         _size = other._size;
-        _capacity = 0;
+        _capacity = other._capacity;
         _shrinkCounter = 0;
+        _owned = false;
     }
 }
 
 IceInternal::Buffer::Container::~Container()
 {
-    if(_buf && _capacity > 0)
+    if(_buf && _owned)
     {
         ::free(_buf);
     }
@@ -91,25 +99,27 @@ IceInternal::Buffer::Container::swap(Container& other)
     std::swap(_size, other._size);
     std::swap(_capacity, other._capacity);
     std::swap(_shrinkCounter, other._shrinkCounter);
+    std::swap(_owned, other._owned);
 }
 
 void
 IceInternal::Buffer::Container::clear()
 {
-    if(_buf && _capacity > 0)
+    if(_buf && _owned)
     {
-        free(_buf);
+        ::free(_buf);
     }
+
     _buf = 0;
     _size = 0;
     _capacity = 0;
+    _shrinkCounter = 0;
+    _owned = true;
 }
 
 void
 IceInternal::Buffer::Container::reserve(size_type n)
 {
-    assert(!_buf || _capacity > 0);
-
     size_type c = _capacity;
     if(n > _capacity)
     {
@@ -125,11 +135,26 @@ IceInternal::Buffer::Container::reserve(size_type n)
         return;
     }
 
-    pointer p = reinterpret_cast<pointer>(::realloc(_buf, _capacity));
+    pointer p;
+    if(_owned)
+    {
+        p = reinterpret_cast<pointer>(::realloc(_buf, _capacity));
+    }
+    else
+    {
+        p = reinterpret_cast<pointer>(::malloc(_capacity));
+        if(p)
+        {
+            ::memcpy(p, _buf, _size);
+            _owned = true;
+        }
+    }
+
     if(!p)
     {
         _capacity = c; // Restore the previous capacity.
         throw std::bad_alloc();
     }
+
     _buf = p;
 }
diff --git a/cpp/src/Ice/OutputStream.cpp b/cpp/src/Ice/OutputStream.cpp
index 05c800a6e10..c9b4ca0a5de 100644
--- a/cpp/src/Ice/OutputStream.cpp
+++ b/cpp/src/Ice/OutputStream.cpp
@@ -95,6 +95,16 @@ Ice::OutputStream::OutputStream(const CommunicatorPtr& communicator, const Encod
     initialize(communicator, encoding);
 }
 
+Ice::OutputStream::OutputStream(const CommunicatorPtr& communicator, const EncodingVersion& encoding,
+                                const pair<const Byte*, const Byte*>& buf) :
+    Buffer(buf.first, buf.second),
+    _closure(0),
+    _currentEncaps(0)
+{
+    initialize(communicator, encoding);
+    b.reset();
+}
+
 Ice::OutputStream::OutputStream(Instance* instance, const EncodingVersion& encoding) :
     _closure(0),
     _currentEncaps(0)
diff --git a/cpp/src/IceDB/IceDB.h b/cpp/src/IceDB/IceDB.h
index af9bd2206c0..d905f2bab22 100644
--- a/cpp/src/IceDB/IceDB.h
+++ b/cpp/src/IceDB/IceDB.h
@@ -515,19 +515,13 @@ struct Codec<T, IceContext, Ice::OutputStream>
 
     static bool write(const T& t, MDB_val& val, const IceContext& ctx)
     {
-        Ice::OutputStream stream(ctx.communicator, ctx.encoding);
+        const size_t limit = val.mv_size;
+        std::pair<Ice::Byte*, Ice::Byte*> p(reinterpret_cast<Ice::Byte*>(val.mv_data),
+                                            reinterpret_cast<Ice::Byte*>(val.mv_data) + limit);
+        Ice::OutputStream stream(ctx.communicator, ctx.encoding, p);
         stream.write(t);
-        if(stream.b.size() > val.mv_size)
-        {
-            val.mv_size = stream.b.size();
-            return false;
-        }
-        else
-        {
-            val.mv_size = stream.b.size();
-            memcpy(val.mv_data, &stream.b[0], stream.b.size());
-            return true;
-        }
+        val.mv_size = stream.b.size();
+        return stream.b.size() > limit ? false : true;
     }
 };
 
diff --git a/cpp/test/Ice/stream/Client.cpp b/cpp/test/Ice/stream/Client.cpp
index 25e0a165976..418dd020c28 100644
--- a/cpp/test/Ice/stream/Client.cpp
+++ b/cpp/test/Ice/stream/Client.cpp
@@ -1241,6 +1241,39 @@ run(int, char**, const Ice::CommunicatorPtr& communicator)
         }
     }
 
+    //
+    // Test marshaling to user-supplied buffer.
+    //
+    {
+        Ice::Byte buf[128];
+        pair<Ice::Byte*, Ice::Byte*> p(&buf[0], &buf[0] + sizeof(buf));
+        Ice::OutputStream out(communicator, Ice::currentEncoding, p);
+        vector<Ice::Byte> v;
+        v.resize(127);
+        out.write(v);
+        test(out.pos() == 128); // 127 bytes + leading size (1 byte)
+        test(out.b.begin() == buf); // Verify the stream hasn't reallocated.
+    }
+    {
+        Ice::Byte buf[128];
+        pair<Ice::Byte*, Ice::Byte*> p(&buf[0], &buf[0] + sizeof(buf));
+        Ice::OutputStream out(communicator, Ice::currentEncoding, p);
+        vector<Ice::Byte> v;
+        v.resize(127);
+        ::memset(&v[0], 0xFF, v.size());
+        out.write(v);
+        out.write(Ice::Byte(0xFF)); // This extra byte should make the stream reallocate.
+        test(out.pos() == 129); // 127 bytes + leading size (1 byte) + 1 byte
+        test(out.b.begin() != buf); // Verify the stream was reallocated.
+        out.finished(data);
+
+        Ice::InputStream in(communicator, data);
+        vector<Ice::Byte> v2;
+        in.read(v2);
+        test(v2.size() == 127);
+        test(v == v2); // Make sure the original buffer was preserved.
+    }
+
     cout << "ok" << endl;
     return 0;
 }