Merge pull request #86 from thatstoasty/socket

Introduce Socket and refactor connection caching

Author: saviorand

.github/workflows/bench.yml

@@ -0,0 +1,18 @@
+name: Run the benchmarking suite
+
+on:
+  workflow_call:
+
+jobs:
+  test:
+    name: Run benchmarks
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Run the test suite
+        run: |
+          curl -ssL https://magic.modular.com | bash
+          source $HOME/.bash_profile
+          magic run bench
+          # magic run bench_server # Commented out until we get `wrk` installed

.github/workflows/branch.yml

@@ -2,10 +2,10 @@ name: Branch workflow
 
 on:
   push:
-    branches: 
+    branches:
       - '*'
   pull_request:
-    branches: 
+    branches:
       - '*'
 
 permissions:
@@ -14,6 +14,9 @@ permissions:
 jobs:
   test:
     uses: ./.github/workflows/test.yml
-  
+
+  bench:
+    uses: ./.github/workflows/bench.yml
+
   package:
     uses: ./.github/workflows/package.yml

.github/workflows/test.yml

@@ -15,4 +15,5 @@ jobs:
           curl -ssL https://magic.modular.com | bash
           source $HOME/.bash_profile
           magic run test
-
+          magic run integration_tests_py
+          magic run integration_tests_external

.gitignore

@@ -15,4 +15,6 @@ install_id
 output
 
 # misc
-.vscode
+.vscode
+
+__pycache__

bench.mojo renamed to benchmark/bench.mojo

@@ -74,9 +74,8 @@ fn lightbug_benchmark_response_parse(mut b: Bencher):
     @always_inline
     @parameter
     fn response_parse():
-        var res = Response
         try:
-            _ = HTTPResponse.from_bytes(res.as_bytes())
+            _ = HTTPResponse.from_bytes(Response.as_bytes())
         except:
             pass
 
@@ -88,9 +87,8 @@ fn lightbug_benchmark_request_parse(mut b: Bencher):
     @always_inline
     @parameter
     fn request_parse():
-        var r = Request
         try:
-            _ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, r.as_bytes())
+            _ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, Request.as_bytes())
         except:
             pass
 
@@ -103,7 +101,7 @@ fn lightbug_benchmark_request_encode(mut b: Bencher):
     @parameter
     fn request_encode():
         var req = HTTPRequest(
-            URI.parse("http://127.0.0.1:8080/some-path")[URI],
+            URI.parse("http://127.0.0.1:8080/some-path"),
             headers=headers_struct,
             body=body_bytes,
         )
@@ -118,8 +116,7 @@ fn lightbug_benchmark_header_encode(mut b: Bencher):
     @parameter
     fn header_encode():
         var b = ByteWriter()
-        var h = headers_struct
-        b.write(h)
+        b.write(headers_struct)
 
     b.iter[header_encode]()
 
@@ -130,9 +127,8 @@ fn lightbug_benchmark_header_parse(mut b: Bencher):
     @parameter
     fn header_parse():
         try:
-            var b = headers
             var header = Headers()
-            var reader = ByteReader(b.as_bytes())
+            var reader = ByteReader(headers.as_bytes())
             _ = header.parse_raw(reader)
         except:
             print("failed")

bench_server.mojo renamed to benchmark/bench_server.mojo

@@ -3,7 +3,7 @@ from lightbug_http.client import Client
 
 
 fn test_request(mut client: Client) raises -> None:
-    var uri = URI.parse_raises("google.com")
+    var uri = URI.parse("google.com")
     var headers = Headers(Header("Host", "google.com"))
     var request = HTTPRequest(uri, headers)
     var response = client.do(request^)

client.mojo

@@ -5,8 +5,3 @@ from lightbug_http.cookie import Cookie, RequestCookieJar, ResponseCookieJar
 from lightbug_http.service import HTTPService, Welcome, Counter
 from lightbug_http.server import Server
 from lightbug_http.strings import to_string
-
-
-trait DefaultConstructible:
-    fn __init__(out self) raises:
-        ...

lightbug_http/__init__.mojo

@@ -1,4 +1,6 @@
-from .libc import (
+from collections import Dict
+from memory import UnsafePointer
+from lightbug_http.libc import (
     c_int,
     AF_INET,
     SOCK_STREAM,
@@ -12,32 +14,32 @@ from lightbug_http.strings import to_string
 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, SysConnection
+from lightbug_http.net import create_connection, TCPConnection
 from lightbug_http.io.bytes import Bytes
 from lightbug_http.utils import ByteReader, logger
-from collections import Dict
+from lightbug_http.pool_manager import PoolManager
 
 
 struct Client:
     var host: String
     var port: Int
     var name: String
+    var allow_redirects: Bool
 
-    var _connections: Dict[String, SysConnection]
+    var _connections: PoolManager[TCPConnection]
 
-    fn __init__(out self, host: String = "127.0.0.1", port: Int = 8888):
+    fn __init__(
+        out self,
+        host: String = "127.0.0.1",
+        port: Int = 8888,
+        cached_connections: Int = 10,
+        allow_redirects: Bool = False,
+    ):
         self.host = host
         self.port = port
         self.name = "lightbug_http_client"
-        self._connections = Dict[String, SysConnection]()
-
-    fn __del__(owned self):
-        for conn in self._connections.values():
-            try:
-                conn[].close()
-            except:
-                # TODO: Add an optional debug log entry here
-                pass
+        self.allow_redirects = allow_redirects
+        self._connections = PoolManager[TCPConnection](cached_connections)
 
     fn do(mut self, owned req: HTTPRequest) raises -> HTTPResponse:
         """The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
@@ -84,17 +86,15 @@ struct Client:
             else:
                 port = 80
 
-        var conn: SysConnection
         var cached_connection = False
+        var conn: TCPConnection
         try:
-            conn = self._connections[host_str]
+            conn = self._connections.take(host_str)
             cached_connection = True
-        except:
-            # If connection is not cached, create a new one.
-            try:
-                conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
-                self._connections[host_str] = conn
-            except e:
+        except e:
+            if str(e) == "PoolManager.take: Key not found.":
+                conn = create_connection(host_str, port)
+            else:
                 logger.error(e)
                 raise Error("Client.do: Failed to create a connection to host.")
 
@@ -105,35 +105,49 @@ struct Client:
             # 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.")
-                self._close_conn(host_str)
+                conn.teardown()
                 if cached_connection:
                     return self.do(req^)
             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.
+        # 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)
-        var bytes_recv = conn.read(new_buf)
 
-        if bytes_recv == 0:
-            self._close_conn(host_str)
-            if cached_connection:
-                return self.do(req^)
-            raise Error("Client.do: No response received from the server.")
+        try:
+            _ = conn.read(new_buf)
+        except e:
+            if str(e) == "EOF":
+                conn.teardown()
+                if cached_connection:
+                    return self.do(req^)
+                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
         try:
-            var res = HTTPResponse.from_bytes(new_buf, conn)
-            if res.is_redirect():
-                self._close_conn(host_str)
-                return self._handle_redirect(req^, res^)
-            if res.connection_close():
-                self._close_conn(host_str)
-            return res
+            res = HTTPResponse.from_bytes(new_buf, conn)
         except e:
-            self._close_conn(host_str)
+            logger.error("Failed to parse a response...")
+            try:
+                conn.teardown()
+            except:
+                logger.error("Failed to teardown connection...")
             raise e
 
-        return HTTPResponse(Bytes())
+        # 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():
+            conn.teardown()
+            return self._handle_redirect(req^, res^)
+        # Server told the client to close the connection, we can assume the server closed their side after sending the response.
+        elif res.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
 
     fn _handle_redirect(
         mut self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
@@ -144,20 +158,12 @@ struct Client:
             new_location = original_response.headers[HeaderKey.LOCATION]
         except e:
             raise Error("Client._handle_redirect: `Location` header was not received in the response.")
-        
+
         if new_location and new_location.startswith("http"):
-            try:
-                new_uri = URI.parse_raises(new_location)
-            except e:
-                raise Error("Client._handle_redirect: Failed to parse the new URI - " + str(e))
+            new_uri = URI.parse(new_location)
             original_req.headers[HeaderKey.HOST] = new_uri.host
         else:
             new_uri = original_req.uri
             new_uri.path = new_location
         original_req.uri = new_uri
         return self.do(original_req^)
-
-    fn _close_conn(mut self, host: String) raises:
-        if host in self._connections:
-            self._connections[host].close()
-            _ = self._connections.pop(host)

lightbug_http/client.mojo

@@ -86,7 +86,7 @@ struct Cookie(CollectionElement):
         self.partitioned = partitioned
 
     fn __str__(self) -> String:
-        return "Name: " + self.name + " Value: " + self.value
+        return String.write("Name: ", self.name, " Value: ", self.value)
 
     fn __copyinit__(out self: Cookie, existing: Cookie):
         self.name = existing.name
@@ -101,15 +101,15 @@ struct Cookie(CollectionElement):
         self.partitioned = existing.partitioned
 
     fn __moveinit__(out self: Cookie, owned existing: Cookie):
-        self.name = existing.name
-        self.value = existing.value
-        self.max_age = existing.max_age
-        self.expires = existing.expires
-        self.domain = existing.domain
-        self.path = existing.path
+        self.name = existing.name^
+        self.value = existing.value^
+        self.max_age = existing.max_age^
+        self.expires = existing.expires^
+        self.domain = existing.domain^
+        self.path = existing.path^
         self.secure = existing.secure
         self.http_only = existing.http_only
-        self.same_site = existing.same_site
+        self.same_site = existing.same_site^
         self.partitioned = existing.partitioned
 
     fn clear_cookie(mut self):
@@ -120,23 +120,23 @@ struct Cookie(CollectionElement):
         return Header(HeaderKey.SET_COOKIE, self.build_header_value())
 
     fn build_header_value(self) -> String:
-        var header_value = self.name + Cookie.EQUAL + self.value
+        var header_value = String.write(self.name, Cookie.EQUAL, self.value)
         if self.expires.is_datetime():
             var v = self.expires.http_date_timestamp()
             if v:
-                header_value += Cookie.SEPERATOR + Cookie.EXPIRES + Cookie.EQUAL + v.value()
+                header_value.write(Cookie.SEPERATOR, Cookie.EXPIRES, Cookie.EQUAL, v.value())
         if self.max_age:
-            header_value += Cookie.SEPERATOR + Cookie.MAX_AGE + Cookie.EQUAL + str(self.max_age.value().total_seconds)
+            header_value.write(Cookie.SEPERATOR, Cookie.MAX_AGE, Cookie.EQUAL, str(self.max_age.value().total_seconds))
         if self.domain:
-            header_value += Cookie.SEPERATOR + Cookie.DOMAIN + Cookie.EQUAL + self.domain.value()
+            header_value.write(Cookie.SEPERATOR, Cookie.DOMAIN, Cookie.EQUAL, self.domain.value())
         if self.path:
-            header_value += Cookie.SEPERATOR + Cookie.PATH + Cookie.EQUAL + self.path.value()
+            header_value.write(Cookie.SEPERATOR, Cookie.PATH, Cookie.EQUAL, self.path.value())
         if self.secure:
-            header_value += Cookie.SEPERATOR + Cookie.SECURE
+            header_value.write(Cookie.SEPERATOR, Cookie.SECURE)
         if self.http_only:
-            header_value += Cookie.SEPERATOR + Cookie.HTTP_ONLY
+            header_value.write(Cookie.SEPERATOR, Cookie.HTTP_ONLY)
         if self.same_site:
-            header_value += Cookie.SEPERATOR + Cookie.SAME_SITE + Cookie.EQUAL + str(self.same_site.value())
+            header_value.write(Cookie.SEPERATOR, Cookie.SAME_SITE, Cookie.EQUAL, str(self.same_site.value()))
         if self.partitioned:
-            header_value += Cookie.SEPERATOR + Cookie.PARTITIONED
+            header_value.write(Cookie.SEPERATOR, Cookie.PARTITIONED)
         return header_value

lightbug_http/cookie/cookie.mojo

@@ -3,6 +3,7 @@ from small_time import SmallTime
 alias HTTP_DATE_FORMAT = "ddd, DD MMM YYYY HH:mm:ss ZZZ"
 alias TZ_GMT = TimeZone(0, "GMT")
 
+
 @value
 struct Expiration(CollectionElement):
     var variant: UInt8

lightbug_http/cookie/expiration.mojo

@@ -1,7 +1,6 @@
 from lightbug_http.http import HTTPResponse
-from lightbug_http.io.bytes import bytes
 
-alias TODO_MESSAGE = String("TODO").as_bytes()
+alias TODO_MESSAGE = "TODO".as_bytes()
 
 
 # TODO: Custom error handlers provided by the user