Merge pull request #69 from bgreni/persistent-client-connection

Keep persistent connections in client

Author: saviorand

lightbug_http/client.mojo

@@ -1,4 +1,4 @@
-from lightbug_http.libc import (
+from .libc import (
     c_int,
     AF_INET,
     SOCK_STREAM,
@@ -9,29 +9,43 @@ from lightbug_http.libc import (
     close,
 )
 from lightbug_http.strings import to_string
-from lightbug_http.io.bytes import Bytes
-from lightbug_http.utils import ByteReader
-from lightbug_http.net import create_connection, default_buffer_size
+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.io.bytes import Bytes
+from lightbug_http.utils import ByteReader
+from collections import Dict
 
 
 struct Client:
     var host: StringLiteral
     var port: Int
     var name: String
 
+    var _connections: Dict[String, SysConnection]
+
     fn __init__(inout self) raises:
         self.host = "127.0.0.1"
         self.port = 8888
         self.name = "lightbug_http_client"
+        self._connections = Dict[String, SysConnection]()
 
     fn __init__(inout self, host: StringLiteral, port: Int) raises:
         self.host = host
         self.port = port
         self.name = "lightbug_http_client"
+        self._connections = Dict[String, SysConnection]()
 
-    fn do(self, owned req: HTTPRequest) raises -> HTTPResponse:
+    fn __del__(owned self):
+        for conn in self._connections.values():
+            try:
+                conn[].close()
+            except:
+                # TODO: Add an optional debug log entry here
+                pass
+
+    fn do(inout self, owned req: HTTPRequest) raises -> HTTPResponse:
         """
         The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
 
@@ -83,31 +97,47 @@ struct Client:
             else:
                 port = 80
 
-        # TODO: Actually handle persistent connections
-        var conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
+        var conn: SysConnection
+        var cached_connection = False
+        if host_str in self._connections:
+            conn = self._connections[host_str]
+            cached_connection = True
+        else:
+            conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
+            self._connections[host_str] = conn
+
         var bytes_sent = conn.write(encode(req))
         if bytes_sent == -1:
+            # Maybe peer reset ungracefully, so try a fresh connection
+            self._close_conn(host_str)
+            if cached_connection:
+                return self.do(req^)
             raise Error("Failed to send message")
-        
+
         var new_buf = Bytes(capacity=default_buffer_size)
         var bytes_recv = conn.read(new_buf)
 
         if bytes_recv == 0:
-            conn.close()
+            self._close_conn(host_str)
+            if cached_connection:
+                return self.do(req^)
+            raise Error("No response received")
         try:
             var res = HTTPResponse.from_bytes(new_buf^)
             if res.is_redirect():
-                conn.close()
+                self._close_conn(host_str)
                 return self._handle_redirect(req^, res^)
+            if res.connection_close():
+                self._close_conn(host_str)
             return res
         except e:
-            conn.close()
+            self._close_conn(host_str)
             raise e
 
         return HTTPResponse(Bytes())
 
     fn _handle_redirect(
-        self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
+        inout self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
     ) raises -> HTTPResponse:
         var new_uri: URI
         var new_location = original_response.headers[HeaderKey.LOCATION]
@@ -119,3 +149,7 @@ struct Client:
             new_uri.path = new_location
         original_req.uri = new_uri
         return self.do(original_req^)
+
+    fn _close_conn(inout self, host: String) raises:
+        self._connections[host].close()
+        _ = self._connections.pop(host)

lightbug_http/header.mojo

@@ -15,6 +15,7 @@ struct HeaderKey:
     alias DATE = "date"
     alias LOCATION = "location"
     alias HOST = "host"
+    alias SERVER = "server"
 
 
 @value

lightbug_http/http/request.mojo

@@ -15,6 +15,7 @@ from lightbug_http.strings import (
     to_string,
 )
 
+
 @value
 struct HTTPRequest(Formattable, Stringable):
     var headers: Headers

lightbug_http/http/response.mojo

@@ -125,16 +125,10 @@ struct HTTPResponse(Formattable, Stringable):
         self.set_content_length(len(self.body_raw))
 
     fn format_to(self, inout writer: Formatter):
-        writer.write(
-            self.protocol,
-            whitespace,
-            self.status_code,
-            whitespace,
-            self.status_text,
-            lineBreak,
-            "server: lightbug_http",
-            lineBreak,
-        )
+        writer.write(self.protocol, whitespace, self.status_code, whitespace, self.status_text, lineBreak)
+
+        if HeaderKey.SERVER not in self.headers:
+            writer.write("server: lightbug_http", lineBreak)
 
         self.headers.format_to(writer)
 

magic.lock

@@ -12,7 +12,6 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-h4bc722e_7.conda
       - conda: https://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2024.8.30-hbcca054_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.1.7-unix_pyh707e725_0.conda
-      - conda: https://repo.prefix.dev/mojo-community/linux-64/gojo-0.1.9-hb0f4dca_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.5.0-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/importlib_metadata-8.5.0-hd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_client-8.6.2-pyhd8ed1ab_0.conda
@@ -71,7 +70,6 @@ environments:
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-h99b78c6_7.conda
       - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ca-certificates-2024.8.30-hf0a4a13_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.1.7-unix_pyh707e725_0.conda
-      - conda: https://repo.prefix.dev/mojo-community/osx-arm64/gojo-0.1.9-h60d57d3_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.5.0-pyha770c72_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/importlib_metadata-8.5.0-hd8ed1ab_0.conda
       - conda: https://conda.anaconda.org/conda-forge/noarch/jupyter_client-8.6.2-pyhd8ed1ab_0.conda
@@ -216,34 +214,6 @@ packages:
   license_family: BSD
   size: 84437
   timestamp: 1692311973840
-- kind: conda
-  name: gojo
-  version: 0.1.9
-  build: h60d57d3_0
-  subdir: osx-arm64
-  url: https://repo.prefix.dev/mojo-community/osx-arm64/gojo-0.1.9-h60d57d3_0.conda
-  sha256: 4c268d0d8d5f1b78a547e78a3db9e8037143918cd1d696f5adb5db55942cef5e
-  depends:
-  - max >=24.5.0,<24.6.0
-  arch: arm64
-  platform: osx
-  license: MIT
-  size: 1009999
-  timestamp: 1726268309700
-- kind: conda
-  name: gojo
-  version: 0.1.9
-  build: hb0f4dca_0
-  subdir: linux-64
-  url: https://repo.prefix.dev/mojo-community/linux-64/gojo-0.1.9-hb0f4dca_0.conda
-  sha256: 9a49e21b4269368a6d906769bd041b8b91b99da3375d7944f7d8ddd73392c2f0
-  depends:
-  - max >=24.5.0,<24.6.0
-  arch: x86_64
-  platform: linux
-  license: MIT
-  size: 1011206
-  timestamp: 1726268249824
 - kind: conda
   name: importlib-metadata
   version: 8.5.0

mojoproject.toml

@@ -16,5 +16,4 @@ format = { cmd = "magic run mojo format -l 120 lightbug_http" }
 
 [dependencies]
 max = ">=24.5.0,<25"
-gojo = "0.1.9"
 small_time = "0.1.3"

recipes/recipe.yaml

@@ -19,7 +19,6 @@ build:
 requirements:
   run:
     - max >=24.5.0
-    - gojo == 0.1.9
     - small_time == 0.1.3
 
 about: