test: clean up test-tcp-writealot.c

Avoid unnecessary heap allocations, makes it easier to reason about correctness.

test/test-tcp-writealot.c

@@ -31,10 +31,8 @@
 
 #define TOTAL_BYTES       (WRITES * CHUNKS_PER_WRITE * CHUNK_SIZE)
 
-
 static char* send_buffer;
 
-
 static int shutdown_cb_called = 0;
 static int connect_cb_called = 0;
 static int write_cb_called = 0;
@@ -43,20 +41,18 @@ static int bytes_sent = 0;
 static int bytes_sent_done = 0;
 static int bytes_received_done = 0;
 
+static uv_connect_t connect_req;
+static uv_shutdown_t shutdown_req;
+static uv_write_t write_reqs[WRITES];
+
 
 static uv_buf_t alloc_cb(uv_handle_t* handle, size_t size) {
-  uv_buf_t buf;
-  buf.base = (char*)malloc(size);
-  buf.len = size;
-  return buf;
+  return uv_buf_init(malloc(size), size);
 }
 
 
 static void close_cb(uv_handle_t* handle) {
   ASSERT(handle != NULL);
-
-  free(handle);
-
   close_cb_called++;
 }
 
@@ -64,7 +60,7 @@ static void close_cb(uv_handle_t* handle) {
 static void shutdown_cb(uv_shutdown_t* req, int status) {
   uv_tcp_t* tcp;
 
-  ASSERT(req);
+  ASSERT(req == &shutdown_req);
   ASSERT(status == 0);
 
   tcp = (uv_tcp_t*)(req->handle);
@@ -77,28 +73,21 @@ static void shutdown_cb(uv_shutdown_t* req, int status) {
 
   /* We should have had all the writes called already. */
   ASSERT(write_cb_called == WRITES);
-
-  free(req);
 }
 
 
 static void read_cb(uv_stream_t* tcp, ssize_t nread, uv_buf_t buf) {
   ASSERT(tcp != NULL);
 
-  if (nread < 0) {
+  if (nread >= 0) {
+    bytes_received_done += nread;
+  }
+  else {
     ASSERT(uv_last_error(uv_default_loop()).code == UV_EOF);
     printf("GOT EOF\n");
-
-    if (buf.base) {
-      free(buf.base);
-    }
-
     uv_close((uv_handle_t*)tcp, close_cb);
-    return;
   }
 
-  bytes_received_done += nread;
-
   free(buf.base);
 }
 
@@ -114,71 +103,55 @@ static void write_cb(uv_write_t* req, int status) {
 
   bytes_sent_done += CHUNKS_PER_WRITE * CHUNK_SIZE;
   write_cb_called++;
-
-  free(req);
 }
 
 
 static void connect_cb(uv_connect_t* req, int status) {
   uv_buf_t send_bufs[CHUNKS_PER_WRITE];
-  uv_tcp_t* tcp;
-  uv_write_t* write_req;
-  uv_shutdown_t* shutdown_req;
+  uv_stream_t* stream;
   int i, j, r;
 
-  ASSERT(req != NULL);
+  ASSERT(req == &connect_req);
   ASSERT(status == 0);
 
-  tcp = (uv_tcp_t*)req->handle;
-
+  stream = req->handle;
   connect_cb_called++;
-  free(req);
 
   /* Write a lot of data */
   for (i = 0; i < WRITES; i++) {
+    uv_write_t* write_req = write_reqs + i;
+
     for (j = 0; j < CHUNKS_PER_WRITE; j++) {
-      send_bufs[j].len = CHUNK_SIZE;
-      send_bufs[j].base = send_buffer + bytes_sent;
+      send_bufs[j] = uv_buf_init(send_buffer + bytes_sent, CHUNK_SIZE);
       bytes_sent += CHUNK_SIZE;
     }
 
-    write_req = malloc(sizeof(uv_write_t));
-    ASSERT(write_req != NULL);
-
-    r = uv_write(write_req, (uv_stream_t*) tcp, (uv_buf_t*)&send_bufs,
-        CHUNKS_PER_WRITE, write_cb);
+    r = uv_write(write_req, stream, send_bufs, CHUNKS_PER_WRITE, write_cb);
     ASSERT(r == 0);
   }
 
-  /* Shutdown on drain. FIXME: dealloc req? */
-  shutdown_req = malloc(sizeof(uv_shutdown_t));
-  ASSERT(shutdown_req != NULL);
-  r = uv_shutdown(shutdown_req, (uv_stream_t*)tcp, shutdown_cb);
+  /* Shutdown on drain. */
+  r = uv_shutdown(&shutdown_req, stream, shutdown_cb);
   ASSERT(r == 0);
 
   /* Start reading */
-  r = uv_read_start((uv_stream_t*)tcp, alloc_cb, read_cb);
+  r = uv_read_start(stream, alloc_cb, read_cb);
   ASSERT(r == 0);
 }
 
 
 TEST_IMPL(tcp_writealot) {
   struct sockaddr_in addr = uv_ip4_addr("127.0.0.1", TEST_PORT);
-  uv_tcp_t* client = (uv_tcp_t*)malloc(sizeof *client);
-  uv_connect_t* connect_req = malloc(sizeof(uv_connect_t));
+  uv_tcp_t client;
   int r;
 
-  ASSERT(client != NULL);
-  ASSERT(connect_req != NULL);
-
-  send_buffer = (char*)malloc(TOTAL_BYTES + 1);
-
+  send_buffer = malloc(TOTAL_BYTES);
   ASSERT(send_buffer != NULL);
 
-  r = uv_tcp_init(uv_default_loop(), client);
+  r = uv_tcp_init(uv_default_loop(), &client);
   ASSERT(r == 0);
 
-  r = uv_tcp_connect(connect_req, client, addr, connect_cb);
+  r = uv_tcp_connect(&connect_req, &client, addr, connect_cb);
   ASSERT(r == 0);
 
   uv_run(uv_default_loop());
@@ -191,5 +164,7 @@ TEST_IMPL(tcp_writealot) {
   ASSERT(bytes_sent_done == TOTAL_BYTES);
   ASSERT(bytes_received_done == TOTAL_BYTES);
 
+  free(send_buffer);
+
   return 0;
 }