baraba/Baradb
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: rate limiting, shared DB instance, TCP_NODELAY fix

Browse Source 
- Add rate limiter to TCP and HTTP servers
- Share single LSMTree between TCP and HTTP servers
- Replace unsafe cast[string]/cast[seq[byte]] with safe helpers
- Stabilize gossip UDP socket with exponential backoff
- Fix TCP_NODELAY via low-level setSockOptInt(IPPROTO_TCP)
  (asyncnet.setSockOpt(OptNoDelay) uses SOL_SOCKET and fails)
- Apply TCP_NODELAY to server and websocket sockets
This commit is contained in:
dimgigov
2026-05-18 19:33:40 +03:00
parent 7e6a45e6b7
commit 8afa998516
13 changed files with 240 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/barabadb/core/gossip.nim
+24 -9
View File
@@ -320,21 +320,36 @@ proc startGossipListener*(gp: GossipProtocol) {.async.} =
gp.sock.setSockOpt(OptReuseAddr, true)
gp.sock.bindAddr(Port(gp.gossipPort))
gp.running = true
var consecutiveErrors = 0
const maxConsecutiveErrors = 10
const baseRetryDelayMs = 100
while gp.running:
try:
let (data, senderAddr, _) = await gp.sock.recvFrom(65535)
consecutiveErrors = 0 # Reset on success
if data.len > 0:
gp.handleIncomingGossip(data, senderAddr)
except:
# Small sleep on error to avoid spin
gp.sock.close()
# Recreate socket for next iteration
try:
gp.sock = newAsyncSocket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)
gp.sock.setSockOpt(OptReuseAddr, true)
gp.sock.bindAddr(Port(gp.gossipPort))
except:
except CatchableError as e:
if not gp.running:
break
inc consecutiveErrors
if consecutiveErrors >= maxConsecutiveErrors:
# Recreate socket after too many errors
try:
gp.sock.close()
except:
discard
try:
gp.sock = newAsyncSocket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)
gp.sock.setSockOpt(OptReuseAddr, true)
gp.sock.bindAddr(Port(gp.gossipPort))
consecutiveErrors = 0
except:
break
# Exponential backoff with cap
let delayMs = min(baseRetryDelayMs * (1 shl min(consecutiveErrors, 6)), 5000)
await sleepAsync(delayMs)
proc startGossip*(gp: GossipProtocol, healthIntervalMs: int = 1000,
gossipIntervalMs: int = 2000) =
src/barabadb/core/httpserver.nim
+17 -3
View File
@@ -19,6 +19,7 @@ import ../protocol/wire
import ../core/websocket
import jwt as jwtlib
import ../protocol/auth
import ../protocol/ratelimit
type
HttpServer* = ref object
@@ -29,6 +30,7 @@ type
metrics*: Metrics
secretKey*: string
authManager*: AuthManager
rateLimiter*: RateLimiter
ws*: WsServer
Metrics* = ref object
@@ -38,13 +40,12 @@ type
selectCount*: int
activeConnections*: int
proc newHttpServer*(config: BaraConfig): HttpServer =
let dataDir = config.dataDir / "server"
let db = newLSMTree(dataDir)
proc newHttpServerWithDb*(config: BaraConfig, db: LSMTree): HttpServer =
let ctx = newExecutionContext(db)
ctx.txnManager = newTxnManager()
let secret = config.getEffectiveJwtSecret()
let ws = newWsServer(config, secret)
let rl = newRateLimiter(rlaTokenBucket, config.rateLimitGlobal, config.rateLimitPerClient)
ctx.onChange = proc(ev: ChangeEvent) =
let msg = $ev.kind & " " & ev.table
asyncCheck ws.broadcastToTable(ev.table, msg)
@@ -52,8 +53,14 @@ proc newHttpServer*(config: BaraConfig): HttpServer =
HttpServer(config: config, running: false, db: db, ctx: ctx,
secretKey: secret,
authManager: am,
rateLimiter: rl,
metrics: Metrics(), ws: ws)
proc newHttpServer*(config: BaraConfig): HttpServer =
let dataDir = config.dataDir / "server"
let db = newLSMTree(dataDir)
return newHttpServerWithDb(config, db)
# ----------------------------------------------------------------------
# JWT helpers
# ----------------------------------------------------------------------
@@ -124,6 +131,13 @@ proc queryHandler(server: HttpServer): RequestHandler =
userId = uid
role = r
# Rate limiting
let clientKey = request.headers["X-Forwarded-For"]
let rateLimitKey = if clientKey.len > 0: clientKey else: "http-global"
if not server.rateLimiter.allowRequest(rateLimitKey):
ctx.json(%*{"error": "Rate limit exceeded"}, 429)
return
let body = parseJson(request.body)
if body == nil or "query" notin body:
ctx.json(%*{"error": "Missing 'query' in request body"}, 400)
src/barabadb/core/server.nim
+61 -20
View File
@@ -8,6 +8,11 @@ import std/os
import std/endians
import std/monotimes
import std/locks
import std/nativesockets
when defined(windows):
from std/winlean import TCP_NODELAY
else:
from std/posix import TCP_NODELAY
import config
import logging
import ../protocol/wire
@@ -22,6 +27,7 @@ import ../core/disttxn
import ../core/replication
import ../core/sharding
import ../core/gossip
import ../protocol/ratelimit
import jwt as jwtlib
type
@@ -37,12 +43,11 @@ type
clusterMembership*: ClusterMembership
gossipProtocol*: GossipProtocol
tls*: TLSContext
rateLimiter*: RateLimiter
activeConnections*: int
activeConnectionsLock*: Lock
proc newServer*(config: BaraConfig): Server =
let dataDir = config.dataDir / "server"
let db = newLSMTree(dataDir)
proc newServerWithDb*(config: BaraConfig, db: LSMTree): Server =
let ctx = newExecutionContext(db)
ctx.txnManager = newTxnManager()
var tls: TLSContext = nil
@@ -85,19 +90,40 @@ proc newServer*(config: BaraConfig): Server =
gp.onSuspect = proc(nodeId: string) {.gcsafe.} =
cm.onNodeSuspect(nodeId)
# Initialize rate limiter
let rl = newRateLimiter(rlaTokenBucket, config.rateLimitGlobal, config.rateLimitPerClient)
result = Server(config: config, running: false, db: db, ctx: ctx,
txnManager: ctx.txnManager, distTxnManager: newDistTxnManager(),
replicationManager: newReplicationManager(),
shardRouter: shardRouter,
clusterMembership: cm,
gossipProtocol: gp,
tls: tls)
tls: tls,
rateLimiter: rl)
initLock(result.activeConnectionsLock)
proc newServer*(config: BaraConfig): Server =
let dataDir = config.dataDir / "server"
let db = newLSMTree(dataDir)
return newServerWithDb(config, db)
# ----------------------------------------------------------------------
# Wire Protocol Helpers
# ----------------------------------------------------------------------
proc bytesToString(data: seq[byte]): string =
## Safely convert seq[byte] to string (copies data)
result = newString(data.len)
for i in 0..<data.len:
result[i] = char(data[i])
proc stringToBytes(data: string): seq[byte] =
## Safely convert string to seq[byte] (copies data)
result = newSeq[byte](data.len)
for i in 0..<data.len:
result[i] = byte(data[i])
proc readUint32BE(data: string, pos: int): uint32 =
if pos + 4 > data.len:
raise newException(ValueError, "readUint32BE: index out of bounds")
@@ -381,7 +407,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
let key = data[0..<nullPos]
let value = data[nullPos+1..^1]
if value.len > 0:
server.db.put(key, cast[seq[byte]](value))
server.db.put(key, stringToBytes(value))
else:
server.db.delete(key)
await client.send("ACK " & $lsn & "\n")
@@ -435,27 +461,35 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
if not authenticated and header.kind != mkAuth:
let err = makeErrorMessage(header.requestId, 401, "Authentication required")
await client.send(cast[string](err))
await client.send(bytesToString(err))
continue
# Rate limiting for query messages
if header.kind in {mkQuery, mkQueryParams}:
let clientKey = "client-" & $clientId
if not server.rateLimiter.allowRequest(clientKey):
let err = serializeError(429, "Rate limit exceeded", header.requestId)
await client.send(bytesToString(err))
continue
case header.kind
of mkAuth:
let tokenStr = parseAuthMessage(cast[seq[byte]](payload))
let tokenStr = parseAuthMessage(stringToBytes(payload))
let (valid, userId, role) = verifyToken(secret, tokenStr)
if valid:
authenticated = true
connCtx.currentUser = userId
connCtx.currentRole = role
let okMsg = makeAuthOkMessage(header.requestId)
await client.send(cast[string](okMsg))
await client.send(bytesToString(okMsg))
info("Client " & $clientId & " authenticated as " & userId)
else:
let err = makeErrorMessage(header.requestId, 403, "Invalid token")
await client.send(cast[string](err))
await client.send(bytesToString(err))
of mkQuery:
var pos = 0
let queryStr = readString(cast[seq[byte]](payload), pos)
let queryStr = readString(stringToBytes(payload), pos)
info("[" & $clientId & "] Query: " & queryStr)
# Shard-aware routing: check if this node should handle the write
@@ -471,7 +505,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
if localShards.len == 0:
shardCheck = false
let err = serializeError(3, "Node not assigned to any shard", header.requestId)
await client.send(cast[string](err))
await client.send(bytesToString(err))
break
if shardCheck:
@@ -484,15 +518,15 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
if success:
let dataMsg = serializeResult(result, header.requestId)
await client.send(cast[string](dataMsg))
await client.send(bytesToString(dataMsg))
let completeMsg = serializeComplete(result.affectedRows, header.requestId)
await client.send(cast[string](completeMsg))
await client.send(bytesToString(completeMsg))
else:
let errorMsg = serializeError(1, errorMsg, header.requestId)
await client.send(cast[string](errorMsg))
await client.send(bytesToString(errorMsg))
of mkQueryParams:
let (queryStr, params) = readQueryParamsMessage(cast[seq[byte]](payload))
let (queryStr, params) = readQueryParamsMessage(stringToBytes(payload))
info("[" & $clientId & "] QueryParams: " & queryStr & " (" & $params.len & " params)")
let startTicks = getMonoTime().ticks()
@@ -504,26 +538,26 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
if success:
let dataMsg = serializeResult(result, header.requestId)
await client.send(cast[string](dataMsg))
await client.send(bytesToString(dataMsg))
let completeMsg = serializeComplete(result.affectedRows, header.requestId)
await client.send(cast[string](completeMsg))
await client.send(bytesToString(completeMsg))
else:
let errorMsg = serializeError(1, errorMsg, header.requestId)
await client.send(cast[string](errorMsg))
await client.send(bytesToString(errorMsg))
of mkPing:
var pongMsg = WireMessage(
header: MessageHeader(kind: mkPong, length: 0, requestId: header.requestId),
payload: @[],
)
await client.send(cast[string](serializeMessage(pongMsg)))
await client.send(bytesToString(serializeMessage(pongMsg)))
of mkClose:
break
else:
let errorMsg = serializeError(2, "Unsupported message kind: " & $header.kind, header.requestId)
await client.send(cast[string](errorMsg))
await client.send(bytesToString(errorMsg))
except Exception as e:
errorMsg("Client " & $clientId & " error: " & e.msg)
@@ -537,6 +571,11 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.}
info("Client " & $clientId & " disconnected")
client.close()
proc setTcpNoDelay(sock: AsyncSocket) =
## Enable TCP_NODELAY using the correct protocol level (IPPROTO_TCP).
## asyncnet.setSockOpt uses SOL_SOCKET by default which fails for OptNoDelay.
setSockOptInt(sock.getFd, IPPROTO_TCP.cint, TCP_NODELAY, 1)
proc run*(server: Server) {.async.} =
server.running = true
var clientId = 0
@@ -548,6 +587,7 @@ proc run*(server: Server) {.async.} =
let sock = newAsyncSocket()
sock.setSockOpt(OptReuseAddr, true)
setTcpNoDelay(sock)
sock.bindAddr(Port(server.config.port), server.config.address)
sock.listen()
if server.config.tlsEnabled:
@@ -556,6 +596,7 @@ proc run*(server: Server) {.async.} =
info("BaraDB listening on " & server.config.address & ":" & $server.config.port)
while server.running:
let client = await sock.accept()
setTcpNoDelay(client)
acquire(server.activeConnectionsLock)
var shouldAccept = true
try:
src/barabadb/core/websocket.nim
+11
View File
@@ -5,6 +5,11 @@ import std/strutils
import std/tables
import std/base64
import std/sets
import std/nativesockets
when defined(windows):
from std/winlean import TCP_NODELAY
else:
from std/posix import TCP_NODELAY
import config
import jwt as jwtlib
@@ -307,16 +312,22 @@ proc handleConnection(server: WsServer, client: AsyncSocket) {.async.} =
inc server.nextId
asyncCheck server.handleWsClient(client, server.nextId)
proc setTcpNoDelay(sock: AsyncSocket) =
## Enable TCP_NODELAY using the correct protocol level (IPPROTO_TCP).
setSockOptInt(sock.getFd, IPPROTO_TCP.cint, TCP_NODELAY, 1)
proc run*(server: WsServer, port: int = 9471) {.async.} =
server.running = true
let sock = newAsyncSocket()
sock.setSockOpt(OptReuseAddr, true)
setTcpNoDelay(sock)
sock.bindAddr(Port(port))
sock.listen()
echo "BaraDB WebSocket listening on port ", port
while server.running:
let client = await sock.accept()
setTcpNoDelay(client)
asyncCheck server.handleConnection(client)
proc stop*(server: WsServer) =
src/baradadb.nim
+9 -7
View File
@@ -230,17 +230,21 @@ proc main() =
warn("Failed to generate self-signed certificate. TLS disabled.")
config.tlsEnabled = false
# Shared LSMTree instance for both TCP and HTTP servers
let dataDir = config.dataDir / "server"
let sharedDb = newLSMTree(dataDir)
# Start HTTP server (blocking, multi-threaded via hunos) in background thread
var httpServer = newHttpServer(config)
var httpServer = newHttpServerWithDb(config, sharedDb)
spawn httpServer.run(config.port + 440) # HTTP port = TCP port + 440
# Start background compaction loop
let cm = newCompactionManager(httpServer.db)
let cm = newCompactionManager(sharedDb)
asyncCheck cm.startCompactionLoop()
# Create TCP server (initialization is synchronous, run is async)
let localId {.used.} = if config.raftNodeId.len > 0: config.raftNodeId else: "node-" & $config.port
var tcpServer = newServer(config)
var tcpServer = newServerWithDb(config, sharedDb)
# Start Raft cluster if enabled
if config.raftEnabled:
@@ -251,11 +255,9 @@ proc main() =
if cmd == "put":
let parts = cast[string](data).split("\x00")
if parts.len >= 2:
httpServer.db.put(parts[0], cast[seq[byte]](parts[1]))
tcpServer.db.put(parts[0], cast[seq[byte]](parts[1]))
sharedDb.put(parts[0], cast[seq[byte]](parts[1]))
elif cmd == "delete":
httpServer.db.delete(cast[string](data))
tcpServer.db.delete(cast[string](data))
sharedDb.delete(cast[string](data))
# Wire RAFT ↔ DistTxn
wireRaftDistTxn(raftNode, tcpServer)