[libc] Use RAII based alloc in gpu rpc_server instead of manual new/delete (llvm#110341)

Co-authored-by: Joseph Huber <[email protected]>

Author: Hardcode84

libc/utils/gpu/server/rpc_server.cpp

@@ -206,6 +206,17 @@ void handle_printf(rpc::Server::Port &port) {
     delete[] reinterpret_cast<char *>(ptr);
 }
 
+namespace {
+struct TempStorage {
+  char *alloc(size_t size) {
+    storage.emplace_back(std::make_unique<char[]>(size));
+    return storage.back().get();
+  }
+
+  std::vector<std::unique_ptr<char[]>> storage;
+};
+} // namespace
+
 template <uint32_t lane_size>
 rpc_status_t handle_server_impl(
     rpc::Server &server,
@@ -216,6 +227,8 @@ rpc_status_t handle_server_impl(
   if (!port)
     return RPC_STATUS_SUCCESS;
 
+  TempStorage temp_storage;
+
   switch (port->get_opcode()) {
   case RPC_WRITE_TO_STREAM:
   case RPC_WRITE_TO_STDERR:
@@ -234,29 +247,28 @@ rpc_status_t handle_server_impl(
       std::fill(files, files + lane_size, stdout);
     }
 
-    port->recv_n(strs, sizes, [&](uint64_t size) { return new char[size]; });
+    port->recv_n(strs, sizes,
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->send([&](rpc::Buffer *buffer, uint32_t id) {
       flockfile(files[id]);
       buffer->data[0] = fwrite_unlocked(strs[id], 1, sizes[id], files[id]);
       if (port->get_opcode() == RPC_WRITE_TO_STDOUT_NEWLINE &&
           buffer->data[0] == sizes[id])
         buffer->data[0] += fwrite_unlocked("\n", 1, 1, files[id]);
       funlockfile(files[id]);
-      delete[] reinterpret_cast<uint8_t *>(strs[id]);
     });
     break;
   }
   case RPC_READ_FROM_STREAM: {
     uint64_t sizes[lane_size] = {0};
     void *data[lane_size] = {nullptr};
     port->recv([&](rpc::Buffer *buffer, uint32_t id) {
-      data[id] = new char[buffer->data[0]];
+      data[id] = temp_storage.alloc(buffer->data[0]);
       sizes[id] =
           fread(data[id], 1, buffer->data[0], file::to_stream(buffer->data[1]));
     });
     port->send_n(data, sizes);
     port->send([&](rpc::Buffer *buffer, uint32_t id) {
-      delete[] reinterpret_cast<uint8_t *>(data[id]);
       std::memcpy(buffer->data, &sizes[id], sizeof(uint64_t));
     });
     break;
@@ -265,27 +277,24 @@ rpc_status_t handle_server_impl(
     uint64_t sizes[lane_size] = {0};
     void *data[lane_size] = {nullptr};
     port->recv([&](rpc::Buffer *buffer, uint32_t id) {
-      data[id] = new char[buffer->data[0]];
+      data[id] = temp_storage.alloc(buffer->data[0]);
       const char *str =
           fgets(reinterpret_cast<char *>(data[id]), buffer->data[0],
                 file::to_stream(buffer->data[1]));
       sizes[id] = !str ? 0 : std::strlen(str) + 1;
     });
     port->send_n(data, sizes);
-    for (uint32_t id = 0; id < lane_size; ++id)
-      if (data[id])
-        delete[] reinterpret_cast<uint8_t *>(data[id]);
     break;
   }
   case RPC_OPEN_FILE: {
     uint64_t sizes[lane_size] = {0};
     void *paths[lane_size] = {nullptr};
-    port->recv_n(paths, sizes, [&](uint64_t size) { return new char[size]; });
+    port->recv_n(paths, sizes,
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->recv_and_send([&](rpc::Buffer *buffer, uint32_t id) {
       FILE *file = fopen(reinterpret_cast<char *>(paths[id]),
                          reinterpret_cast<char *>(buffer->data));
       buffer->data[0] = reinterpret_cast<uintptr_t>(file);
-      delete[] reinterpret_cast<uint8_t *>(paths[id]);
     });
     break;
   }
@@ -316,13 +325,12 @@ rpc_status_t handle_server_impl(
   case RPC_HOST_CALL: {
     uint64_t sizes[lane_size] = {0};
     void *args[lane_size] = {nullptr};
-    port->recv_n(args, sizes, [&](uint64_t size) { return new char[size]; });
+    port->recv_n(args, sizes,
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->recv([&](rpc::Buffer *buffer, uint32_t id) {
       reinterpret_cast<void (*)(void *)>(buffer->data[0])(args[id]);
     });
-    port->send([&](rpc::Buffer *, uint32_t id) {
-      delete[] reinterpret_cast<uint8_t *>(args[id]);
-    });
+    port->send([&](rpc::Buffer *, uint32_t id) {});
     break;
   }
   case RPC_FEOF: {
@@ -385,11 +393,11 @@ rpc_status_t handle_server_impl(
   case RPC_REMOVE: {
     uint64_t sizes[lane_size] = {0};
     void *args[lane_size] = {nullptr};
-    port->recv_n(args, sizes, [&](uint64_t size) { return new char[size]; });
+    port->recv_n(args, sizes,
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->send([&](rpc::Buffer *buffer, uint32_t id) {
       buffer->data[0] = static_cast<uint64_t>(
           remove(reinterpret_cast<const char *>(args[id])));
-      delete[] reinterpret_cast<uint8_t *>(args[id]);
     });
     break;
   }
@@ -399,26 +407,24 @@ rpc_status_t handle_server_impl(
     void *oldpath[lane_size] = {nullptr};
     void *newpath[lane_size] = {nullptr};
     port->recv_n(oldpath, oldsizes,
-                 [&](uint64_t size) { return new char[size]; });
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->recv_n(newpath, newsizes,
-                 [&](uint64_t size) { return new char[size]; });
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->send([&](rpc::Buffer *buffer, uint32_t id) {
       buffer->data[0] = static_cast<uint64_t>(
           rename(reinterpret_cast<const char *>(oldpath[id]),
                  reinterpret_cast<const char *>(newpath[id])));
-      delete[] reinterpret_cast<uint8_t *>(oldpath[id]);
-      delete[] reinterpret_cast<uint8_t *>(newpath[id]);
     });
     break;
   }
   case RPC_SYSTEM: {
     uint64_t sizes[lane_size] = {0};
     void *args[lane_size] = {nullptr};
-    port->recv_n(args, sizes, [&](uint64_t size) { return new char[size]; });
+    port->recv_n(args, sizes,
+                 [&](uint64_t size) { return temp_storage.alloc(size); });
     port->send([&](rpc::Buffer *buffer, uint32_t id) {
       buffer->data[0] = static_cast<uint64_t>(
           system(reinterpret_cast<const char *>(args[id])));
-      delete[] reinterpret_cast<uint8_t *>(args[id]);
     });
     break;
   }