Add global allocator to avoid unnecessary allocations / copies when c…

…alling preview2 ABI

expected/wasm32-wasip2/defined-symbols.txt

@@ -1429,6 +1429,10 @@ vwprintf
 vwscanf
 wall_clock_now
 wall_clock_resolution
+wasip2_allocator_alloc
+wasip2_allocator_is_buffer
+wasip2_allocator_reset
+wasip2_allocator_set
 wasip2_list_string_free
 wasip2_list_tuple2_string_string_free
 wasip2_list_u32_free

libc-bottom-half/headers/private/wasi/wasip2_allocator.h

@@ -0,0 +1,14 @@
+#ifndef WASIP2_ALLOCATOR_H
+#define WASIP2_ALLOCATOR_H
+
+#include <inttypes.h>
+
+void wasip2_allocator_reset();
+
+void wasip2_allocator_set(void *buffer, size_t len);
+
+void* wasip2_allocator_alloc(size_t len);
+
+int wasip2_allocator_is_buffer(void *buf);
+
+#endif

libc-bottom-half/sources/recv.c

@@ -6,6 +6,7 @@
 #include <wasi/api.h>
 #include <wasi/descriptor_table.h>
 #include <wasi/sockets_utils.h>
+#include <wasi/wasip2_allocator.h>
 
 static ssize_t tcp_recvfrom(tcp_socket_t *socket, uint8_t *buffer,
 			    size_t length, int flags, struct sockaddr *addr,
@@ -42,6 +43,9 @@ static ssize_t tcp_recvfrom(tcp_socket_t *socket, uint8_t *buffer,
 
 	streams_borrow_input_stream_t rx_borrow =
 		streams_borrow_input_stream(connection.input);
+
+	wasip2_allocator_set(buffer, length);
+
 	while (true) {
 		wasip2_list_u8_t result;
 		streams_stream_error_t error;
@@ -52,20 +56,25 @@ static ssize_t tcp_recvfrom(tcp_socket_t *socket, uint8_t *buffer,
 			} else {
 				// TODO wasi-sockets: wasi-sockets has no way to recover TCP stream errors yet.
 				errno = EPIPE;
+				wasip2_allocator_reset();
 				return -1;
 			}
 		}
 
 		if (result.len) {
-			memcpy(buffer, result.ptr, result.len);
-			wasip2_list_u8_free(&result);
+			if (!wasip2_allocator_is_buffer(result.ptr)) {
+				memcpy(buffer, result.ptr, result.len);
+				wasip2_list_u8_free(&result);
+			}
+			wasip2_allocator_reset();
 			return result.len;
 		} else if (should_block) {
 			poll_borrow_pollable_t pollable_borrow =
 				poll_borrow_pollable(connection.input_pollable);
 			poll_method_pollable_block(pollable_borrow);
 		} else {
 			errno = EWOULDBLOCK;
+			wasip2_allocator_reset();
 			return -1;
 		}
 	}

libc-bottom-half/sources/wasip2.c

@@ -1,6 +1,6 @@
 // Generated by `wit-bindgen` 0.17.0. DO NOT EDIT!
 #include "wasi/wasip2.h"
-
+#include "wasi/wasip2_allocator.h"
 
 __attribute__((__import_module__("wasi:cli/[email protected]"), __import_name__("get-environment")))
 extern void __wasm_import_environment_get_environment(int32_t);
@@ -368,11 +368,12 @@ extern int64_t __wasm_import_random_insecure_get_insecure_random_u64(void);
 __attribute__((__import_module__("wasi:random/[email protected]"), __import_name__("insecure-seed")))
 extern void __wasm_import_random_insecure_seed_insecure_seed(int32_t);
 
-__attribute__((__weak__, __export_name__("cabi_realloc")))
+__attribute__((__weak__))
 void *cabi_realloc(void *ptr, size_t old_size, size_t align, size_t new_size) {
   (void) old_size;
+  if (ptr) abort(); // not supported yet
   if (new_size == 0) return (void*) align;
-  void *ret = realloc(ptr, new_size);
+  void *ret = wasip2_allocator_alloc(new_size);
   if (!ret) abort();
   return ret;
 }
@@ -1128,7 +1129,7 @@ bool streams_method_input_stream_read(streams_borrow_input_stream_t self, uint64
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1167,7 +1168,7 @@ bool streams_method_input_stream_blocking_read(streams_borrow_input_stream_t sel
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1206,7 +1207,7 @@ bool streams_method_input_stream_skip(streams_borrow_input_stream_t self, uint64
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1245,7 +1246,7 @@ bool streams_method_input_stream_blocking_skip(streams_borrow_input_stream_t sel
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1289,7 +1290,7 @@ bool streams_method_output_stream_check_write(streams_borrow_output_stream_t sel
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1327,7 +1328,7 @@ bool streams_method_output_stream_write(streams_borrow_output_stream_t self, was
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1364,7 +1365,7 @@ bool streams_method_output_stream_blocking_write_and_flush(streams_borrow_output
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1401,7 +1402,7 @@ bool streams_method_output_stream_flush(streams_borrow_output_stream_t self, str
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1438,7 +1439,7 @@ bool streams_method_output_stream_blocking_flush(streams_borrow_output_stream_t
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1480,7 +1481,7 @@ bool streams_method_output_stream_write_zeroes(streams_borrow_output_stream_t se
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1517,7 +1518,7 @@ bool streams_method_output_stream_blocking_write_zeroes_and_flush(streams_borrow
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1555,7 +1556,7 @@ bool streams_method_output_stream_splice(streams_borrow_output_stream_t self, st
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -1594,7 +1595,7 @@ bool streams_method_output_stream_blocking_splice(streams_borrow_output_stream_t
           break;
         }
       }
-      
+
       result.val.err = variant;
       break;
     }
@@ -2199,7 +2200,7 @@ bool filesystem_method_descriptor_stat(filesystem_borrow_descriptor_t self, file
           break;
         }
       }
-      
+
       result.val.ok = (filesystem_descriptor_stat_t) {
         (int32_t) (*((uint8_t*) (ptr + 8))),
         (uint64_t) (*((int64_t*) (ptr + 16))),
@@ -2279,7 +2280,7 @@ bool filesystem_method_descriptor_stat_at(filesystem_borrow_descriptor_t self, f
           break;
         }
       }
-      
+
       result.val.ok = (filesystem_descriptor_stat_t) {
         (int32_t) (*((uint8_t*) (ptr + 8))),
         (uint64_t) (*((int64_t*) (ptr + 16))),
@@ -2646,7 +2647,7 @@ bool filesystem_method_directory_entry_stream_read_directory_entry(filesystem_bo
           break;
         }
       }
-      
+
       result.val.ok = option;
       break;
     }
@@ -2961,7 +2962,7 @@ bool udp_method_udp_socket_local_address(udp_borrow_udp_socket_t self, udp_ip_so
           break;
         }
       }
-      
+
       result.val.ok = variant;
       break;
     }
@@ -3023,7 +3024,7 @@ bool udp_method_udp_socket_remote_address(udp_borrow_udp_socket_t self, udp_ip_s
           break;
         }
       }
-      
+
       result.val.ok = variant;
       break;
     }
@@ -3647,7 +3648,7 @@ bool tcp_method_tcp_socket_local_address(tcp_borrow_tcp_socket_t self, tcp_ip_so
           break;
         }
       }
-      
+
       result.val.ok = variant;
       break;
     }
@@ -3709,7 +3710,7 @@ bool tcp_method_tcp_socket_remote_address(tcp_borrow_tcp_socket_t self, tcp_ip_s
           break;
         }
       }
-      
+
       result.val.ok = variant;
       break;
     }
@@ -4254,12 +4255,12 @@ bool ip_name_lookup_method_resolve_address_stream_resolve_next_address(ip_name_l
               break;
             }
           }
-          
+
           option.val = variant;
           break;
         }
       }
-      
+
       result.val.ok = option;
       break;
     }

libc-bottom-half/sources/wasip2_allocator.c

@@ -0,0 +1,38 @@
+#include <wasi/wasip2_allocator.h>
+#include <stdlib.h>
+#include <string.h>
+
+typedef struct {
+    void* buffer;
+    size_t len;
+    int used;
+} wasip2_allocator_t;
+
+static _Thread_local wasip2_allocator_t _allocator = {0};
+
+void wasip2_allocator_set(void *buffer, size_t len)
+{
+    _allocator.buffer = buffer;
+    _allocator.len = len;
+    _allocator.used = 0;
+}
+
+void wasip2_allocator_reset()
+{
+    memset(&_allocator, 0, sizeof(_allocator));
+}
+
+void* wasip2_allocator_alloc(size_t len)
+{
+    if (_allocator.used == 0 && len <= _allocator.len) {
+        void *ret = _allocator.buffer;
+        _allocator.used = 1;
+        return ret;
+    }
+    return malloc(len);
+}
+
+int wasip2_allocator_is_buffer(void *buf)
+{
+    return buf == _allocator.buffer;
+}