fix socket bug, enable __del__ and __enter__, and switch poolmanager to use tuple of host, port, and scheme

Author: thatstoasty

lightbug_http/client.mojo

@@ -1,12 +1,35 @@
 from collections import Dict
+from utils import StringSlice
 from memory import UnsafePointer
 from lightbug_http.net import default_buffer_size
 from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.header import Headers, HeaderKey
 from lightbug_http.net import create_connection, TCPConnection
 from lightbug_http.io.bytes import Bytes
 from lightbug_http.utils import ByteReader, logger
-from lightbug_http.pool_manager import PoolManager
+from lightbug_http.pool_manager import PoolManager, Scheme, PoolKey
+
+
+fn parse_host_and_port(source: String, is_tls: Bool) raises -> (String, UInt16):
+    """Parses the host and port from a given string.
+
+    Args:
+        source: The host uri to parse.
+        is_tls: A boolean indicating whether the connection is secure.
+
+    Returns:
+        A tuple containing the host and port.
+    """
+    var port: UInt16
+    if source.count(":") != 1:
+        port = 443 if is_tls else 80
+        return source, port
+
+    var host: String
+    var reader = ByteReader(source.as_bytes())
+    host = StringSlice(unsafe_from_utf8=reader.read_until(ord(":")))
+    port = atol(StringSlice(unsafe_from_utf8=reader.read_bytes()[1:]))
+    return host^, port
 
 
 struct Client:
@@ -30,7 +53,7 @@ struct Client:
         self.allow_redirects = allow_redirects
         self._connections = PoolManager[TCPConnection](cached_connections)
 
-    fn do(mut self, owned req: HTTPRequest) raises -> HTTPResponse:
+    fn do(mut self, owned request: HTTPRequest) raises -> HTTPResponse:
         """The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
 
         It performs the following steps:
@@ -43,81 +66,67 @@ struct Client:
         Note: The code assumes that the `HTTPRequest` object passed as an argument has a valid URI with a host and port specified.
 
         Args:
-            req: An `HTTPRequest` object representing the request to be sent.
+            request: An `HTTPRequest` object representing the request to be sent.
 
         Returns:
             The received response.
 
         Raises:
             Error: If there is a failure in sending or receiving the message.
         """
-        if req.uri.host == "":
+        if request.uri.host == "":
             raise Error("Client.do: Request failed because the host field is empty.")
-        var is_tls = False
 
-        if req.uri.is_https():
+        var is_tls = False
+        var scheme = Scheme.HTTP
+        if request.uri.is_https():
             is_tls = True
+            scheme = Scheme.HTTPS
 
-        var host_str: String
-        var port: Int
-        if ":" in req.uri.host:
-            var host_port: List[String]
-            try:
-                host_port = req.uri.host.split(":")
-            except:
-                raise Error("Client.do: Failed to split host and port.")
-            host_str = host_port[0]
-            port = atol(host_port[1])
-        else:
-            host_str = req.uri.host
-            if is_tls:
-                port = 443
-            else:
-                port = 80
-
+        host, port = parse_host_and_port(request.uri.host, is_tls)
+        var pool_key = PoolKey(host, port, scheme)
         var cached_connection = False
         var conn: TCPConnection
         try:
-            conn = self._connections.take(host_str)
+            conn = self._connections.take(pool_key)
             cached_connection = True
         except e:
             if str(e) == "PoolManager.take: Key not found.":
-                conn = create_connection(host_str, port)
+                conn = create_connection(host, port)
             else:
                 logger.error(e)
                 raise Error("Client.do: Failed to create a connection to host.")
 
         var bytes_sent: Int
         try:
-            bytes_sent = conn.write(encode(req))
+            bytes_sent = conn.write(encode(request))
         except e:
             # Maybe peer reset ungracefully, so try a fresh connection
             if str(e) == "SendError: Connection reset by peer.":
                 logger.debug("Client.do: Connection reset by peer. Trying a fresh connection.")
                 conn.teardown()
                 if cached_connection:
-                    return self.do(req^)
+                    return self.do(request^)
             logger.error("Client.do: Failed to send message.")
             raise e
 
         # TODO: What if the response is too large for the buffer? We should read until the end of the response. (@thatstoasty)
         var new_buf = Bytes(capacity=default_buffer_size)
-
         try:
             _ = conn.read(new_buf)
         except e:
             if str(e) == "EOF":
                 conn.teardown()
                 if cached_connection:
-                    return self.do(req^)
+                    return self.do(request^)
                 raise Error("Client.do: No response received from the server.")
             else:
                 logger.error(e)
                 raise Error("Client.do: Failed to read response from peer.")
 
-        var res: HTTPResponse
+        var response: HTTPResponse
         try:
-            res = HTTPResponse.from_bytes(new_buf, conn)
+            response = HTTPResponse.from_bytes(new_buf, conn)
         except e:
             logger.error("Failed to parse a response...")
             try:
@@ -127,19 +136,19 @@ struct Client:
             raise e
 
         # Redirects should not keep the connection alive, as redirects can send the client to a different server.
-        if self.allow_redirects and res.is_redirect():
+        if self.allow_redirects and response.is_redirect():
             conn.teardown()
-            return self._handle_redirect(req^, res^)
+            return self._handle_redirect(request^, response^)
         # Server told the client to close the connection, we can assume the server closed their side after sending the response.
-        elif res.connection_close():
+        elif response.connection_close():
             conn.teardown()
         # Otherwise, persist the connection by giving it back to the pool manager.
         else:
-            self._connections.give(host_str, conn^)
-        return res
+            self._connections.give(pool_key, conn^)
+        return response
 
     fn _handle_redirect(
-        mut self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
+        mut self, owned original_request: HTTPRequest, owned original_response: HTTPResponse
     ) raises -> HTTPResponse:
         var new_uri: URI
         var new_location: String
@@ -150,9 +159,9 @@ struct Client:
 
         if new_location and new_location.startswith("http"):
             new_uri = URI.parse(new_location)
-            original_req.headers[HeaderKey.HOST] = new_uri.host
+            original_request.headers[HeaderKey.HOST] = new_uri.host
         else:
-            new_uri = original_req.uri
+            new_uri = original_request.uri
             new_uri.path = new_location
-        original_req.uri = new_uri
-        return self.do(original_req^)
+        original_request.uri = new_uri
+        return self.do(original_request^)

lightbug_http/pool_manager.mojo

@@ -9,40 +9,100 @@ from lightbug_http.utils import logger
 from lightbug_http.owning_list import OwningList
 
 
+@value
+struct Scheme(Hashable, EqualityComparable, Representable, Stringable, Writable):
+    var value: String
+    alias HTTP = Self("http")
+    alias HTTPS = Self("https")
+
+    fn __hash__(self) -> UInt:
+        return hash(self.value)
+
+    fn __eq__(self, other: Self) -> Bool:
+        return self.value == other.value
+
+    fn __ne__(self, other: Self) -> Bool:
+        return self.value != other.value
+
+    fn write_to[W: Writer, //](self, mut writer: W) -> None:
+        writer.write("Scheme(value=", repr(self.value), ")")
+
+    fn __repr__(self) -> String:
+        return String.write(self)
+
+    fn __str__(self) -> String:
+        return self.value.upper()
+
+
+@value
+struct PoolKey(Hashable, KeyElement):
+    var host: String
+    var port: UInt16
+    var scheme: Scheme
+
+    fn __init__(out self, host: String, port: UInt16, scheme: Scheme):
+        self.host = host
+        self.port = port
+        self.scheme = scheme
+
+    fn __hash__(self) -> UInt:
+        # TODO: Very rudimentary hash. We probably need to actually have an actual hash function here.
+        # Since Tuple doesn't have one.
+        return hash(self.host) + hash(self.port) + hash(self.scheme)
+
+    fn __eq__(self, other: Self) -> Bool:
+        return self.host == other.host and self.port == other.port and self.scheme == other.scheme
+
+    fn __ne__(self, other: Self) -> Bool:
+        return self.host != other.host or self.port != other.port or self.scheme != other.scheme
+
+    fn __str__(self) -> String:
+        var result = String()
+        result.write(self.scheme.value, "://", self.host, ":", str(self.port))
+        return result
+
+    fn __repr__(self) -> String:
+        return String.write(self)
+
+    fn write_to[W: Writer, //](self, mut writer: W) -> None:
+        writer.write(
+            "PoolKey(", "scheme=", repr(self.scheme.value), ", host=", repr(self.host), ", port=", str(self.port), ")"
+        )
+
+
 struct PoolManager[ConnectionType: Connection]():
     var _connections: OwningList[ConnectionType]
     var _capacity: Int
-    var mapping: Dict[String, Int]
+    var mapping: Dict[PoolKey, Int]
 
     fn __init__(out self, capacity: Int = 10):
         self._connections = OwningList[ConnectionType](capacity=capacity)
         self._capacity = capacity
-        self.mapping = Dict[String, Int]()
+        self.mapping = Dict[PoolKey, Int]()
 
     fn __del__(owned self):
         logger.debug(
             "PoolManager shutting down and closing remaining connections before destruction:", self._connections.size
         )
         self.clear()
 
-    fn give(mut self, host: String, owned value: ConnectionType) raises:
-        if host in self.mapping:
-            self._connections[self.mapping[host]] = value^
+    fn give(mut self, key: PoolKey, owned value: ConnectionType) raises:
+        if key in self.mapping:
+            self._connections[self.mapping[key]] = value^
             return
 
         if self._connections.size == self._capacity:
             raise Error("PoolManager.give: Cache is full.")
 
-        self._connections[self._connections.size] = value^
-        self.mapping[host] = self._connections.size
-        self._connections.size += 1
-        logger.debug("Checked in connection for peer:", host + ", at index:", self._connections.size)
+        self._connections.append(value^)
+        self.mapping[key] = self._connections.size - 1
+        logger.debug("Checked in connection for peer:", str(key) + ", at index:", self._connections.size)
 
-    fn take(mut self, host: String) raises -> ConnectionType:
+    fn take(mut self, key: PoolKey) raises -> ConnectionType:
         var index: Int
         try:
-            index = self.mapping[host]
-            _ = self.mapping.pop(host)
+            index = self.mapping[key]
+            _ = self.mapping.pop(key)
         except:
             raise Error("PoolManager.take: Key not found.")
 
@@ -52,7 +112,7 @@ struct PoolManager[ConnectionType: Connection]():
             if kv[].value > index:
                 self.mapping[kv[].key] -= 1
 
-        logger.debug("Checked out connection for peer:", host + ", from index:", self._connections.size + 1)
+        logger.debug("Checked out connection for peer:", str(key) + ", from index:", self._connections.size + 1)
         return connection^
 
     fn clear(mut self):
@@ -65,14 +125,14 @@ struct PoolManager[ConnectionType: Connection]():
                 logger.error("Failed to tear down connection. Error:", e)
         self.mapping.clear()
 
-    fn __contains__(self, host: String) -> Bool:
-        return host in self.mapping
+    fn __contains__(self, key: PoolKey) -> Bool:
+        return key in self.mapping
 
-    fn __setitem__(mut self, host: String, owned value: ConnectionType) raises -> None:
-        if host in self.mapping:
-            self._connections[self.mapping[host]] = value^
+    fn __setitem__(mut self, key: PoolKey, owned value: ConnectionType) raises -> None:
+        if key in self.mapping:
+            self._connections[self.mapping[key]] = value^
         else:
-            self.give(host, value^)
+            self.give(key, value^)
 
-    fn __getitem__(self, host: String) raises -> ref [self._connections] ConnectionType:
-        return self._connections[self.mapping[host]]
+    fn __getitem__(self, key: PoolKey) raises -> ref [self._connections] ConnectionType:
+        return self._connections[self.mapping[key]]

lightbug_http/socket.mojo

@@ -148,10 +148,16 @@ struct Socket[AddrType: Addr, address_family: Int = AF_INET](Representable, Stri
         self.fd = existing.fd
         self.socket_type = existing.socket_type
         self.protocol = existing.protocol
+
         self._local_address = existing._local_address^
+        existing._local_address = AddrType()
         self._remote_address = existing._remote_address^
+        existing._remote_address = AddrType()
+
         self._closed = existing._closed
+        existing._closed = True
         self._connected = existing._connected
+        existing._connected = False
 
     fn teardown(mut self) raises:
         """Close the socket and free the file descriptor."""
@@ -169,17 +175,16 @@ struct Socket[AddrType: Addr, address_family: Int = AF_INET](Representable, Stri
                 raise e
 
     # TODO: Removed until we can determine why __del__ bugs out in the client flow, but not server flow?
-    # fn __enter__(owned self) -> Self:
-    #     return self^
+    fn __enter__(owned self) -> Self:
+        return self^
 
     # TODO: Seems to be bugged if this is included. Mojo tries to delete a mystical 0 fd socket that was never initialized?
-    # fn __del__(owned self):
-    #     """Close the socket when the object is deleted."""
-    #     logger.info("In socket del", self)
-    #     try:
-    #         self.teardown()
-    #     except e:
-    #         logger.debug("Socket.__del__: Failed to close socket during deletion:", e)
+    fn __del__(owned self):
+        """Close the socket when the object is deleted."""
+        try:
+            self.teardown()
+        except e:
+            logger.debug("Socket.__del__: Failed to close socket during deletion:", e)
 
     fn __str__(self) -> String:
         return String.write(self)