diff --git a/src/barabadb/core/server.nim b/src/barabadb/core/server.nim index fc9e25c..c307497 100644 --- a/src/barabadb/core/server.nim +++ b/src/barabadb/core/server.nim @@ -250,6 +250,16 @@ proc verifyToken(secret, tokenStr: string): (bool, string, string) = except: return (false, "", "") +proc recvWithTimeout(client: AsyncSocket, size: int, timeoutMs: int): Future[string] {.async.} = + if timeoutMs <= 0: + return await client.recv(size) + let fut = client.recv(size) + let timeoutFut = sleepAsync(timeoutMs) + await fut or timeoutFut + if fut.finished: + return fut.read() + return "" + proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} = info("Client " & $clientId & " connected") var connCtx = cloneForConnection(server.ctx) @@ -270,7 +280,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} # Text-based 2PC protocol: read rest of line var rest = headerData[7..^1] while '\n' notin rest: - let more = await client.recv(1024) + let more = await client.recvWithTimeout(1024, idleTimeout) if more.len == 0: break rest.add(more) let parts = rest.strip().split(" ") @@ -308,7 +318,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} if headerData.len >= 4 and headerData[0..3] == "REP ": var rest = headerData[4..^1] while '\n' notin rest: - let more = await client.recv(1024) + let more = await client.recvWithTimeout(1024, idleTimeout) if more.len == 0: break rest.add(more) let parts = rest.strip().split(" ") @@ -318,7 +328,7 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} if dataLen > 0: var data = "" while data.len < dataLen: - let chunk = await client.recv(dataLen - data.len) + let chunk = await client.recvWithTimeout(dataLen - data.len, idleTimeout) if chunk.len == 0: break data.add(chunk) # Apply replicated data to database @@ -424,7 +434,8 @@ proc handleClient(server: Server, client: AsyncSocket, clientId: int) {.async.} except Exception as e: errorMsg("Client " & $clientId & " error: " & e.msg) finally: - dec server.activeConnections + if server.activeConnections > 0: + dec server.activeConnections info("Client " & $clientId & " disconnected") client.close() diff --git a/src/barabadb/graph/engine.nim b/src/barabadb/graph/engine.nim index 8ba4dd1..ada118e 100644 --- a/src/barabadb/graph/engine.nim +++ b/src/barabadb/graph/engine.nim @@ -70,6 +70,10 @@ proc addEdge*(g: Graph, src, dst: NodeId, label: string = "", weight: float64 = 1.0): EdgeId = acquire(g.lock) defer: release(g.lock) + if src notin g.nodes: + raise newException(KeyError, "Source node does not exist: " & $uint64(src)) + if dst notin g.nodes: + raise newException(KeyError, "Destination node does not exist: " & $uint64(dst)) let id = EdgeId(g.nextEdgeId) inc g.nextEdgeId g.edges[id] = Edge(id: id, src: src, dst: dst, label: label, diff --git a/src/barabadb/protocol/ratelimit.nim b/src/barabadb/protocol/ratelimit.nim index 4fe4a2c..fb6eae6 100644 --- a/src/barabadb/protocol/ratelimit.nim +++ b/src/barabadb/protocol/ratelimit.nim @@ -90,6 +90,26 @@ proc newRateLimiter*(algo: RateLimitAlgo = rlaTokenBucket, proc allowRequest*(rl: RateLimiter, clientId: string): bool = acquire(rl.lock) + # Global rate enforcement + var globalOk = true + if rl.globalRate > 0: + let globalKey = "__global__" + case rl.algo + of rlaTokenBucket: + if globalKey notin rl.buckets: + rl.buckets[globalKey] = newTokenBucket(float64(rl.globalRate), + float64(rl.globalRate) / 60.0) + globalOk = rl.buckets[globalKey].consume() + of rlaSlidingWindow, rlaFixedWindow: + if globalKey notin rl.windows: + rl.windows[globalKey] = newSlidingWindow(60_000_000_000, rl.globalRate) + globalOk = rl.windows[globalKey].allow() + + if not globalOk: + release(rl.lock) + return false + + # Per-client rate enforcement case rl.algo of rlaTokenBucket: if clientId notin rl.buckets: @@ -122,6 +142,25 @@ proc remainingQuota*(rl: RateLimiter, clientId: string): int = result = rl.perClientRate release(rl.lock) +proc cleanupStaleClients*(rl: RateLimiter, maxClients: int = 10000) = + acquire(rl.lock) + # Remove oldest clients if exceeding max + if rl.buckets.len > maxClients: + var toRemove = rl.buckets.len - maxClients + for key in rl.buckets.keys: + if key != "__global__": + rl.buckets.del(key) + dec toRemove + if toRemove <= 0: break + if rl.windows.len > maxClients: + var toRemove = rl.windows.len - maxClients + for key in rl.windows.keys: + if key != "__global__": + rl.windows.del(key) + dec toRemove + if toRemove <= 0: break + release(rl.lock) + proc resetClient*(rl: RateLimiter, clientId: string) = acquire(rl.lock) rl.buckets.del(clientId) diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index 648bfb1..852c8d7 100644 --- a/src/barabadb/query/executor.nim +++ b/src/barabadb/query/executor.nim @@ -188,7 +188,10 @@ proc exprToSql(node: Node): string = else: " " & $node.binOp & " " return exprToSql(node.binLeft) & opStr & exprToSql(node.binRight) of nkFuncCall: - return node.funcName & "(" & exprToSql(node.funcArgs[0]) & ")" + if node.funcArgs.len > 0: + return node.funcName & "(" & exprToSql(node.funcArgs[0]) & ")" + else: + return node.funcName & "()" of nkUnaryOp: return $node.unOp & " " & exprToSql(node.unOperand) else: @@ -1256,9 +1259,10 @@ proc executePlan*(ctx: ExecutionContext, plan: IRPlan): seq[Row] = if i < plan.projectExprs.len: let expr = plan.projectExprs[i] if expr.kind == irekStar: - for k, v in sourceRows[0]: - if not k.startsWith("$") and not k.contains("."): - newRow[k] = v + if sourceRows.len > 0: + for k, v in sourceRows[0]: + if not k.startsWith("$") and not k.contains("."): + newRow[k] = v elif expr.kind == irekAggregate: case expr.aggOp of irCount: @@ -2004,9 +2008,10 @@ proc executeQuery*(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] # Fire BEFORE INSERT triggers var row = initTable[string, string]() - for i, f in mutableFields: - if i < mutableValues[0].len: - row[f] = mutableValues[0][i] + if mutableValues.len > 0: + for i, f in mutableFields: + if i < mutableValues[0].len: + row[f] = mutableValues[0][i] fireTriggers(ctx, stmt.insTarget, "before", "insert", row) var kvPairs: seq[(string, seq[byte])] diff --git a/src/barabadb/vector/engine.nim b/src/barabadb/vector/engine.nim index 42a310d..3246874 100644 --- a/src/barabadb/vector/engine.nim +++ b/src/barabadb/vector/engine.nim @@ -82,6 +82,8 @@ proc manhattanDistance*(a, b: Vector): float64 = return sum proc distance*(a, b: Vector, metric: DistanceMetric): float64 = + if a.len != b.len: + raise newException(ValueError, "Vector dimension mismatch: " & $a.len & " != " & $b.len) case metric of dmCosine: cosineDistance(a, b) of dmEuclidean: euclideanDistance(a, b)