Skip to content

Commit b41e88e

Browse files
authored
Merge pull request #86 from thatstoasty/socket
Introduce Socket and refactor connection caching
2 parents bcd2967 + 8914f99 commit b41e88e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

57 files changed

+3225
-963
lines changed

.github/workflows/bench.yml

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
name: Run the benchmarking suite
2+
3+
on:
4+
workflow_call:
5+
6+
jobs:
7+
test:
8+
name: Run benchmarks
9+
runs-on: ubuntu-latest
10+
steps:
11+
- name: Checkout code
12+
uses: actions/checkout@v4
13+
- name: Run the test suite
14+
run: |
15+
curl -ssL https://magic.modular.com | bash
16+
source $HOME/.bash_profile
17+
magic run bench
18+
# magic run bench_server # Commented out until we get `wrk` installed

.github/workflows/branch.yml

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@ name: Branch workflow
22

33
on:
44
push:
5-
branches:
5+
branches:
66
- '*'
77
pull_request:
8-
branches:
8+
branches:
99
- '*'
1010

1111
permissions:
@@ -14,6 +14,9 @@ permissions:
1414
jobs:
1515
test:
1616
uses: ./.github/workflows/test.yml
17-
17+
18+
bench:
19+
uses: ./.github/workflows/bench.yml
20+
1821
package:
1922
uses: ./.github/workflows/package.yml

.github/workflows/test.yml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,5 @@ jobs:
1515
curl -ssL https://magic.modular.com | bash
1616
source $HOME/.bash_profile
1717
magic run test
18-
18+
magic run integration_tests_py
19+
magic run integration_tests_external

.gitignore

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,6 @@ install_id
1515
output
1616

1717
# misc
18-
.vscode
18+
.vscode
19+
20+
__pycache__

bench.mojo renamed to benchmark/bench.mojo

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,8 @@ fn lightbug_benchmark_response_parse(mut b: Bencher):
7474
@always_inline
7575
@parameter
7676
fn response_parse():
77-
var res = Response
7877
try:
79-
_ = HTTPResponse.from_bytes(res.as_bytes())
78+
_ = HTTPResponse.from_bytes(Response.as_bytes())
8079
except:
8180
pass
8281

@@ -88,9 +87,8 @@ fn lightbug_benchmark_request_parse(mut b: Bencher):
8887
@always_inline
8988
@parameter
9089
fn request_parse():
91-
var r = Request
9290
try:
93-
_ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, r.as_bytes())
91+
_ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, Request.as_bytes())
9492
except:
9593
pass
9694

@@ -103,7 +101,7 @@ fn lightbug_benchmark_request_encode(mut b: Bencher):
103101
@parameter
104102
fn request_encode():
105103
var req = HTTPRequest(
106-
URI.parse("http://127.0.0.1:8080/some-path")[URI],
104+
URI.parse("http://127.0.0.1:8080/some-path"),
107105
headers=headers_struct,
108106
body=body_bytes,
109107
)
@@ -118,8 +116,7 @@ fn lightbug_benchmark_header_encode(mut b: Bencher):
118116
@parameter
119117
fn header_encode():
120118
var b = ByteWriter()
121-
var h = headers_struct
122-
b.write(h)
119+
b.write(headers_struct)
123120

124121
b.iter[header_encode]()
125122

@@ -130,9 +127,8 @@ fn lightbug_benchmark_header_parse(mut b: Bencher):
130127
@parameter
131128
fn header_parse():
132129
try:
133-
var b = headers
134130
var header = Headers()
135-
var reader = ByteReader(b.as_bytes())
131+
var reader = ByteReader(headers.as_bytes())
136132
_ = header.parse_raw(reader)
137133
except:
138134
print("failed")
File renamed without changes.

client.mojo

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ from lightbug_http.client import Client
33

44

55
fn test_request(mut client: Client) raises -> None:
6-
var uri = URI.parse_raises("google.com")
6+
var uri = URI.parse("google.com")
77
var headers = Headers(Header("Host", "google.com"))
88
var request = HTTPRequest(uri, headers)
99
var response = client.do(request^)

lightbug_http/__init__.mojo

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,3 @@ from lightbug_http.cookie import Cookie, RequestCookieJar, ResponseCookieJar
55
from lightbug_http.service import HTTPService, Welcome, Counter
66
from lightbug_http.server import Server
77
from lightbug_http.strings import to_string
8-
9-
10-
trait DefaultConstructible:
11-
fn __init__(out self) raises:
12-
...

lightbug_http/client.mojo

Lines changed: 55 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
from .libc import (
1+
from collections import Dict
2+
from memory import UnsafePointer
3+
from lightbug_http.libc import (
24
c_int,
35
AF_INET,
46
SOCK_STREAM,
@@ -12,32 +14,32 @@ from lightbug_http.strings import to_string
1214
from lightbug_http.net import default_buffer_size
1315
from lightbug_http.http import HTTPRequest, HTTPResponse, encode
1416
from lightbug_http.header import Headers, HeaderKey
15-
from lightbug_http.net import create_connection, SysConnection
17+
from lightbug_http.net import create_connection, TCPConnection
1618
from lightbug_http.io.bytes import Bytes
1719
from lightbug_http.utils import ByteReader, logger
18-
from collections import Dict
20+
from lightbug_http.pool_manager import PoolManager
1921

2022

2123
struct Client:
2224
var host: String
2325
var port: Int
2426
var name: String
27+
var allow_redirects: Bool
2528

26-
var _connections: Dict[String, SysConnection]
29+
var _connections: PoolManager[TCPConnection]
2730

28-
fn __init__(out self, host: String = "127.0.0.1", port: Int = 8888):
31+
fn __init__(
32+
out self,
33+
host: String = "127.0.0.1",
34+
port: Int = 8888,
35+
cached_connections: Int = 10,
36+
allow_redirects: Bool = False,
37+
):
2938
self.host = host
3039
self.port = port
3140
self.name = "lightbug_http_client"
32-
self._connections = Dict[String, SysConnection]()
33-
34-
fn __del__(owned self):
35-
for conn in self._connections.values():
36-
try:
37-
conn[].close()
38-
except:
39-
# TODO: Add an optional debug log entry here
40-
pass
41+
self.allow_redirects = allow_redirects
42+
self._connections = PoolManager[TCPConnection](cached_connections)
4143

4244
fn do(mut self, owned req: HTTPRequest) raises -> HTTPResponse:
4345
"""The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
@@ -84,17 +86,15 @@ struct Client:
8486
else:
8587
port = 80
8688

87-
var conn: SysConnection
8889
var cached_connection = False
90+
var conn: TCPConnection
8991
try:
90-
conn = self._connections[host_str]
92+
conn = self._connections.take(host_str)
9193
cached_connection = True
92-
except:
93-
# If connection is not cached, create a new one.
94-
try:
95-
conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
96-
self._connections[host_str] = conn
97-
except e:
94+
except e:
95+
if str(e) == "PoolManager.take: Key not found.":
96+
conn = create_connection(host_str, port)
97+
else:
9898
logger.error(e)
9999
raise Error("Client.do: Failed to create a connection to host.")
100100

@@ -105,35 +105,49 @@ struct Client:
105105
# Maybe peer reset ungracefully, so try a fresh connection
106106
if str(e) == "SendError: Connection reset by peer.":
107107
logger.debug("Client.do: Connection reset by peer. Trying a fresh connection.")
108-
self._close_conn(host_str)
108+
conn.teardown()
109109
if cached_connection:
110110
return self.do(req^)
111111
logger.error("Client.do: Failed to send message.")
112112
raise e
113113

114-
# TODO: What if the response is too large for the buffer? We should read until the end of the response.
114+
# TODO: What if the response is too large for the buffer? We should read until the end of the response. (@thatstoasty)
115115
var new_buf = Bytes(capacity=default_buffer_size)
116-
var bytes_recv = conn.read(new_buf)
117116

118-
if bytes_recv == 0:
119-
self._close_conn(host_str)
120-
if cached_connection:
121-
return self.do(req^)
122-
raise Error("Client.do: No response received from the server.")
117+
try:
118+
_ = conn.read(new_buf)
119+
except e:
120+
if str(e) == "EOF":
121+
conn.teardown()
122+
if cached_connection:
123+
return self.do(req^)
124+
raise Error("Client.do: No response received from the server.")
125+
else:
126+
logger.error(e)
127+
raise Error("Client.do: Failed to read response from peer.")
123128

129+
var res: HTTPResponse
124130
try:
125-
var res = HTTPResponse.from_bytes(new_buf, conn)
126-
if res.is_redirect():
127-
self._close_conn(host_str)
128-
return self._handle_redirect(req^, res^)
129-
if res.connection_close():
130-
self._close_conn(host_str)
131-
return res
131+
res = HTTPResponse.from_bytes(new_buf, conn)
132132
except e:
133-
self._close_conn(host_str)
133+
logger.error("Failed to parse a response...")
134+
try:
135+
conn.teardown()
136+
except:
137+
logger.error("Failed to teardown connection...")
134138
raise e
135139

136-
return HTTPResponse(Bytes())
140+
# Redirects should not keep the connection alive, as redirects can send the client to a different server.
141+
if self.allow_redirects and res.is_redirect():
142+
conn.teardown()
143+
return self._handle_redirect(req^, res^)
144+
# Server told the client to close the connection, we can assume the server closed their side after sending the response.
145+
elif res.connection_close():
146+
conn.teardown()
147+
# Otherwise, persist the connection by giving it back to the pool manager.
148+
else:
149+
self._connections.give(host_str, conn^)
150+
return res
137151

138152
fn _handle_redirect(
139153
mut self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
@@ -144,20 +158,12 @@ struct Client:
144158
new_location = original_response.headers[HeaderKey.LOCATION]
145159
except e:
146160
raise Error("Client._handle_redirect: `Location` header was not received in the response.")
147-
161+
148162
if new_location and new_location.startswith("http"):
149-
try:
150-
new_uri = URI.parse_raises(new_location)
151-
except e:
152-
raise Error("Client._handle_redirect: Failed to parse the new URI - " + str(e))
163+
new_uri = URI.parse(new_location)
153164
original_req.headers[HeaderKey.HOST] = new_uri.host
154165
else:
155166
new_uri = original_req.uri
156167
new_uri.path = new_location
157168
original_req.uri = new_uri
158169
return self.do(original_req^)
159-
160-
fn _close_conn(mut self, host: String) raises:
161-
if host in self._connections:
162-
self._connections[host].close()
163-
_ = self._connections.pop(host)

lightbug_http/cookie/cookie.mojo

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ struct Cookie(CollectionElement):
8686
self.partitioned = partitioned
8787

8888
fn __str__(self) -> String:
89-
return "Name: " + self.name + " Value: " + self.value
89+
return String.write("Name: ", self.name, " Value: ", self.value)
9090

9191
fn __copyinit__(out self: Cookie, existing: Cookie):
9292
self.name = existing.name
@@ -101,15 +101,15 @@ struct Cookie(CollectionElement):
101101
self.partitioned = existing.partitioned
102102

103103
fn __moveinit__(out self: Cookie, owned existing: Cookie):
104-
self.name = existing.name
105-
self.value = existing.value
106-
self.max_age = existing.max_age
107-
self.expires = existing.expires
108-
self.domain = existing.domain
109-
self.path = existing.path
104+
self.name = existing.name^
105+
self.value = existing.value^
106+
self.max_age = existing.max_age^
107+
self.expires = existing.expires^
108+
self.domain = existing.domain^
109+
self.path = existing.path^
110110
self.secure = existing.secure
111111
self.http_only = existing.http_only
112-
self.same_site = existing.same_site
112+
self.same_site = existing.same_site^
113113
self.partitioned = existing.partitioned
114114

115115
fn clear_cookie(mut self):
@@ -120,23 +120,23 @@ struct Cookie(CollectionElement):
120120
return Header(HeaderKey.SET_COOKIE, self.build_header_value())
121121

122122
fn build_header_value(self) -> String:
123-
var header_value = self.name + Cookie.EQUAL + self.value
123+
var header_value = String.write(self.name, Cookie.EQUAL, self.value)
124124
if self.expires.is_datetime():
125125
var v = self.expires.http_date_timestamp()
126126
if v:
127-
header_value += Cookie.SEPERATOR + Cookie.EXPIRES + Cookie.EQUAL + v.value()
127+
header_value.write(Cookie.SEPERATOR, Cookie.EXPIRES, Cookie.EQUAL, v.value())
128128
if self.max_age:
129-
header_value += Cookie.SEPERATOR + Cookie.MAX_AGE + Cookie.EQUAL + str(self.max_age.value().total_seconds)
129+
header_value.write(Cookie.SEPERATOR, Cookie.MAX_AGE, Cookie.EQUAL, str(self.max_age.value().total_seconds))
130130
if self.domain:
131-
header_value += Cookie.SEPERATOR + Cookie.DOMAIN + Cookie.EQUAL + self.domain.value()
131+
header_value.write(Cookie.SEPERATOR, Cookie.DOMAIN, Cookie.EQUAL, self.domain.value())
132132
if self.path:
133-
header_value += Cookie.SEPERATOR + Cookie.PATH + Cookie.EQUAL + self.path.value()
133+
header_value.write(Cookie.SEPERATOR, Cookie.PATH, Cookie.EQUAL, self.path.value())
134134
if self.secure:
135-
header_value += Cookie.SEPERATOR + Cookie.SECURE
135+
header_value.write(Cookie.SEPERATOR, Cookie.SECURE)
136136
if self.http_only:
137-
header_value += Cookie.SEPERATOR + Cookie.HTTP_ONLY
137+
header_value.write(Cookie.SEPERATOR, Cookie.HTTP_ONLY)
138138
if self.same_site:
139-
header_value += Cookie.SEPERATOR + Cookie.SAME_SITE + Cookie.EQUAL + str(self.same_site.value())
139+
header_value.write(Cookie.SEPERATOR, Cookie.SAME_SITE, Cookie.EQUAL, str(self.same_site.value()))
140140
if self.partitioned:
141-
header_value += Cookie.SEPERATOR + Cookie.PARTITIONED
141+
header_value.write(Cookie.SEPERATOR, Cookie.PARTITIONED)
142142
return header_value

lightbug_http/cookie/expiration.mojo

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ from small_time import SmallTime
33
alias HTTP_DATE_FORMAT = "ddd, DD MMM YYYY HH:mm:ss ZZZ"
44
alias TZ_GMT = TimeZone(0, "GMT")
55

6+
67
@value
78
struct Expiration(CollectionElement):
89
var variant: UInt8

lightbug_http/error.mojo

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from lightbug_http.http import HTTPResponse
2-
from lightbug_http.io.bytes import bytes
32

4-
alias TODO_MESSAGE = String("TODO").as_bytes()
3+
alias TODO_MESSAGE = "TODO".as_bytes()
54

65

76
# TODO: Custom error handlers provided by the user

0 commit comments

Comments
 (0)