1 files changed, 23 insertions, 1 deletions
diff --git a/lib/port_handlers.cpp b/lib/port_handlers.cpp
index d0d5896..848d9ef 100644
--- a/lib/port_handlers.cpp
+++ b/lib/port_handlers.cpp
@@ -60,6 +60,19 @@ void ElementImpl::handle_input_msg(
         return;
     }
 
+    //handle the upstream block allocation request
+    if (msg.type() == typeid(InputAllocatorMessage))
+    {
+        InputAllocatorMessage message;
+        message.token = block_ptr->input_buffer_allocator(
+            index,
+            msg.cast<InputAllocatorMessage>().token,
+            msg.cast<InputAllocatorMessage>().recommend_length
+        );
+        if (message.token) handle.post_upstream(index, message);
+        return;
+    }
+
     ASSERT(false);
 }
 
@@ -77,7 +90,6 @@ void ElementImpl::handle_output_msg(
         return;
     }
 
-
     //a downstream block has declared itself done, recheck the token
     if (msg.type() == typeid(CheckTokensMessage))
     {
@@ -111,5 +123,15 @@ void ElementImpl::handle_output_msg(
         return;
     }
 
+    //return of a positive downstream allocation
+    //reset the token, and clear old output buffers
+    //the new token from the downstream is installed
+    if (msg.type() == typeid(InputAllocatorMessage))
+    {
+        this->output_buffer_tokens[index].reset();
+        this->output_queues.flush(index);
+        this->output_buffer_tokens[index] = msg.cast<InputAllocatorMessage>().token;
+    }
+
     ASSERT(false);
 }