From a2a24d6a64c7b95b234f3bf0ae940b041352b765 Mon Sep 17 00:00:00 2001
From: GuanLuo <41310872+GuanLuo@users.noreply.github.com>
Date: Tue, 20 Aug 2019 10:19:09 -0700
Subject: [PATCH] Fix alloc_fn in grpc/http frontend to match trtserver.h
 description... (#566)

Check memory type and expect repeated call only if byte size is non-zero
---
 src/servers/grpc_server.cc |  5 +++--
 src/servers/http_server.cc | 14 +++++++-------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/servers/grpc_server.cc b/src/servers/grpc_server.cc
index 9f6f0692c1..7993c6726c 100644
--- a/src/servers/grpc_server.cc
+++ b/src/servers/grpc_server.cc
@@ -138,8 +138,9 @@ AsyncResources::ResponseAlloc(
   *buffer = nullptr;
   *buffer_userp = nullptr;
 
-  // Can't allocate for any memory type other than CPU.
-  if (memory_type != TRTSERVER_MEMORY_CPU) {
+  // Can't allocate for any memory type other than CPU. If byte size is 0,
+  // proceed regardless of memory type as no allocation is required.
+  if ((memory_type != TRTSERVER_MEMORY_CPU) && (byte_size != 0)) {
     LOG_VERBOSE(1) << "GRPC allocation failed for type " << memory_type
                    << " for " << tensor_name;
     return nullptr;  // Success
diff --git a/src/servers/http_server.cc b/src/servers/http_server.cc
index 34194fa16b..c32057fbe6 100644
--- a/src/servers/http_server.cc
+++ b/src/servers/http_server.cc
@@ -301,15 +301,15 @@ HTTPAPIServer::ResponseAlloc(
   *buffer = nullptr;
   *buffer_userp = nullptr;
 
-  // Can't allocate for any memory type other than CPU.
-  if (memory_type != TRTSERVER_MEMORY_CPU) {
-    LOG_VERBOSE(1) << "HTTP allocation failed for type " << memory_type
-                   << " for " << tensor_name;
-    return nullptr;  // Success
-  }
-
   // Don't need to do anything if no memory was requested.
   if (byte_size > 0) {
+    // Can't allocate for any memory type other than CPU.
+    if (memory_type != TRTSERVER_MEMORY_CPU) {
+      LOG_VERBOSE(1) << "HTTP allocation failed for type " << memory_type
+                    << " for " << tensor_name;
+      return nullptr;  // Success
+    }
+    
     // Reserve requested space in evbuffer...
     struct evbuffer_iovec output_iovec;
     if (evbuffer_reserve_space(evhttp_buffer, byte_size, &output_iovec, 1) !=