Correctly handle tcp disconnections

crates/lb/tests/net.lua

@@ -79,9 +79,81 @@ describe("tcp", function()
     end)
   end)
 
+  describe("stream", function()
+    test("no concurrent two reads/writes", function()
+      local listener = net.listen_tcp(net.localhost())
+      local client = net.connect_tcp(listener:local_addr())
+      local server = listener()
+      local reader = spawn(function()
+        assert(client:read(1) == nil) -- this should block first, then return nil from disconnection
+      end)
+      spawn(function()
+        assert(not pcall(client.read, client, 1)) -- this should fail, since the first task is still reading
+      end):await()
+      server:shutdown()
+      reader:await()
+    end)
+
+    test("allow concurrent read/write", function()
+      local listener = net.listen_tcp(net.localhost())
+      local client = net.connect_tcp(listener:local_addr())
+      local server = listener()
+      local reader = spawn(function()
+        assert(client:read(1) == nil) -- this should block first, then return nil from disconnection
+      end)
+      spawn(function()
+        client:write("hello") -- should be able to write while the first task is reading
+      end):await()
+      server:shutdown()
+      reader:await()
+    end)
+
+    test("stop reading from disconnected stream", function()
+      local listener = net.listen_tcp(net.localhost())
+      local client = net.connect_tcp(listener:local_addr())
+      local server = listener()
+      local reader = spawn(function()
+        while client:read(4) ~= nil do
+        end
+        assert(client:try_read(4) == nil)
+        assert(client:read_partial(4) == nil)
+        assert(client:read(4) == nil)
+      end)
+      for _ = 1, 10 do
+        assert(server:write("ping") == true)
+      end
+      sleep(100)
+      server:shutdown()
+      server = nil
+      collectgarbage()
+      reader:await()
+    end)
+
+    test("stop writing to disconnected stream", function()
+      local listener = net.listen_tcp(net.localhost())
+      local client = net.connect_tcp(listener:local_addr())
+      local server = listener()
+      local writer = spawn(function()
+        while client:write("pong") do
+        end
+        assert(client:try_write("pong") == nil)
+        assert(client:write_partial("pong") == nil)
+        assert(client:write("pong") == false)
+      end)
+      for _ = 1, 10 do
+        assert(server:read(4) == "pong")
+      end
+      sleep(100)
+      server:shutdown()
+      server = nil
+      collectgarbage()
+      writer:await()
+    end)
+  end)
+
   describe("listener", function()
     test("accept", function()
-      local listener = net.listen_tcp("127.0.0.1", 0)
+      local listener = net.listen_tcp(net.localhost())
       local addr = listener:local_addr()
       local accepted = false
       local client = net.tcp()