Initial GRPC implementation

Author: wasdennnoch

latte/src/main/java/gg/beemo/latte/util/SuspendingRatelimit.kt renamed to latte/src/main/java/gg/beemo/latte/util/Ratelimit.kt

@@ -5,15 +5,44 @@ import kotlinx.coroutines.sync.Semaphore
 import kotlinx.coroutines.sync.withPermit
 import kotlin.time.Duration
 
-class SuspendingRatelimit(private val burst: Int, private val duration: Duration) {
+class Ratelimit(val burst: Int, val duration: Duration) {
     @Volatile
     private var remainingQuota: Int = burst
+    val remaining: Int get() = remainingQuota
 
     @Volatile
     private var resetTimestamp: Long = 0
+    val resetAt: Long get() = resetTimestamp
 
     private val quotaRequestSem = Semaphore(1)
 
+    suspend fun requestQuota() {
+        quotaRequestSem.withPermit {
+            if (remainingQuota <= 0) {
+                tryResetQuota()
+                val waitTime = calculateWaitTime()
+                delay(waitTime)
+            }
+            tryResetQuota()
+            check(remainingQuota > 0)
+            remainingQuota--
+        }
+    }
+
+    fun tryRequestQuota(): Pair<Boolean, Long> {
+        tryResetQuota()
+        if (remainingQuota <= 0) {
+            return false to calculateWaitTime()
+        }
+        check(remainingQuota > 0)
+        remainingQuota--
+        return true to calculateWaitTime()
+    }
+
+    fun addQuota(amount: Int) {
+        remainingQuota += amount
+    }
+
     fun overrideRatelimit(
         remainingQuota: Int,
         resetTimestamp: Long,
@@ -32,29 +61,4 @@ class SuspendingRatelimit(private val burst: Int, private val duration: Duration
             resetTimestamp = System.currentTimeMillis() + duration.inWholeMilliseconds
         }
     }
-
-    suspend fun requestQuota() {
-        quotaRequestSem.withPermit {
-            if (remainingQuota <= 0) {
-                val waitTime = calculateWaitTime()
-                delay(waitTime)
-            }
-            tryResetQuota()
-
-            check(remainingQuota > 0)
-            remainingQuota--
-        }
-    }
-
-    fun tryRequestQuota(): Pair<Boolean, Long?>  {
-        if (remainingQuota <= 0) {
-            val waitTime = calculateWaitTime()
-            return false to waitTime
-        }
-        tryResetQuota()
-
-        check(remainingQuota > 0)
-        remainingQuota--
-        return true to null
-    }
 }

proto/ratelimit.proto

@@ -0,0 +1,29 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "gg.beemo.vanilla";
+option java_outer_classname = "RatelimitProto";
+
+package ratelimit;
+
+service Ratelimit {
+  rpc RequestQuota (RatelimitRequest) returns (RatelimitResponse) {}
+}
+
+message RatelimitRequest {
+  RatelimitType type = 1;
+  fixed64 clientId = 2;
+}
+
+message RatelimitResponse {
+  bool granted = 1;
+  uint32 limit = 2;
+  uint32 remaining = 3;
+  uint64 reset = 4;
+  uint32 resetAfter = 5;
+}
+
+enum RatelimitType {
+  GLOBAL = 0;
+  IDENTIFY = 1;
+}

vanilla/build.gradle.kts

@@ -1,27 +1,60 @@
 plugins {
     application
     java
-    kotlin("jvm") version "1.9.20"
+    kotlin("jvm") version "2.0.21"
+    id("com.google.protobuf") version "0.9.4"
 }
 
 group = "gg.beemo.vanilla"
 version = "1.0.0"
 
 dependencies {
     // Kotlin
-    val kotlinCoroutinesVersion = "1.7.3"
+    val kotlinCoroutinesVersion = "1.9.0"
     implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:$kotlinCoroutinesVersion")
 
     // Beemo shared code
     implementation("gg.beemo.latte:latte")
 
+    // gRPC
+    val grpcVersion = "1.68.0"
+    implementation("io.grpc:grpc-netty-shaded:$grpcVersion")
+    implementation("io.grpc:grpc-protobuf:$grpcVersion")
+    implementation("io.grpc:grpc-kotlin-stub:1.4.1")
+    implementation("com.google.protobuf:protobuf-kotlin:4.28.2")
+
     // Logging
-    val log4jVersion = "2.22.0"
+    val log4jVersion = "2.24.1"
     implementation("org.apache.logging.log4j:log4j-api:$log4jVersion")
     implementation("org.apache.logging.log4j:log4j-core:$log4jVersion")
     implementation("org.apache.logging.log4j:log4j-slf4j2-impl:$log4jVersion")
 }
 
+protobuf {
+    protoc {
+        artifact = "com.google.protobuf:protoc:4.28.2"
+    }
+    plugins {
+        create("grpc") {
+            artifact = "io.grpc:protoc-gen-grpc-java:1.68.0"
+        }
+        create("grpckt") {
+            artifact = "io.grpc:protoc-gen-grpc-kotlin:1.4.1:jdk8@jar"
+        }
+    }
+    generateProtoTasks {
+        all().forEach {
+            it.plugins {
+                create("grpc")
+                create("grpckt")
+            }
+            it.builtins {
+                create("kotlin")
+            }
+        }
+    }
+}
+
 repositories {
     mavenCentral()
 }

vanilla/src/main/java/gg/beemo/vanilla/Config.java

@@ -10,4 +10,6 @@ public class Config {
 
     public static String RABBIT_PASSWORD = "guest";
 
+    public static int GRPC_PORT = 1337;
+
 }

vanilla/src/main/java/gg/beemo/vanilla/GrpcRatelimitService.kt

@@ -0,0 +1,47 @@
+package gg.beemo.vanilla
+
+import gg.beemo.latte.logging.Log
+import gg.beemo.latte.util.Ratelimit
+import java.util.concurrent.ConcurrentHashMap
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.seconds
+
+// TODO https://github.com/grpc/grpc-kotlin/blob/master/examples/server/src/main/kotlin/io/grpc/examples/animals/AnimalsServer.kt
+
+class GrpcRatelimitService : RatelimitGrpcKt.RatelimitCoroutineImplBase() {
+
+    private val log by Log
+    private val globalRatelimits = RatelimitMap(50, 1.seconds)
+    private val identifyRatelimits = RatelimitMap(1, 5.seconds)
+
+    override suspend fun requestQuota(request: RatelimitRequest): RatelimitResponse {
+        val (ratelimitMap, typeString) = when (request.type) {
+            RatelimitType.GLOBAL -> globalRatelimits to "global"
+            RatelimitType.IDENTIFY -> identifyRatelimits to "identify"
+            else -> throw IllegalArgumentException("Unknown ratelimit type ${request.type}")
+        }
+        val ratelimit = ratelimitMap.getClientRatelimit(request.clientId)
+        // TODO Do we want to make a blocking version?
+        val (granted, resetAfter) = ratelimit.tryRequestQuota()
+        log.debug("Got '{}' quota request from clientId {}, was tranted: {}", typeString, request.clientId, granted)
+
+        return ratelimitResponse {
+            this.granted = granted
+            this.limit = ratelimit.burst
+            this.remaining = ratelimit.remaining
+            this.reset = ratelimit.resetAt
+            this.resetAfter = resetAfter.toInt()
+        }
+    }
+
+}
+
+private class RatelimitMap(private val burst: Int, private val duration: Duration) {
+
+    private val limiters = ConcurrentHashMap<Long, Ratelimit>()
+
+    fun getClientRatelimit(clientId: Long): Ratelimit = limiters.computeIfAbsent(clientId) {
+        Ratelimit(burst, duration)
+    }
+
+}

vanilla/src/main/java/gg/beemo/vanilla/KafkaRatelimitClient.kt

@@ -6,19 +6,19 @@ import gg.beemo.latte.broker.IgnoreRpcRequest
 import gg.beemo.latte.broker.rpc.RpcStatus
 import gg.beemo.latte.logging.Log
 import gg.beemo.latte.ratelimit.SharedRatelimitData
-import gg.beemo.latte.util.SuspendingRatelimit
+import gg.beemo.latte.util.Ratelimit
 import java.util.concurrent.ConcurrentHashMap
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 
 // Give request expiry a bit of leeway in case of clock drift
 private val EXPIRY_GRACE_PERIOD = 5.seconds.inWholeMilliseconds
 
-class RatelimitClient(connection: BrokerConnection) : BrokerClient(connection) {
+class KafkaRatelimitClient(connection: BrokerConnection) : BrokerClient(connection) {
 
     private val log by Log
-    private val globalRatelimitProvider = RatelimitProvider(50, 1.seconds)
-    private val identifyRatelimitProvider = RatelimitProvider(1, 5.seconds)
+    private val globalRatelimitProvider = KafkaRatelimitProvider(50, 1.seconds)
+    private val identifyRatelimitProvider = KafkaRatelimitProvider(1, 5.seconds)
 
     init {
         rpc<SharedRatelimitData.RatelimitRequestData, Unit>(
@@ -54,12 +54,12 @@ class RatelimitClient(connection: BrokerConnection) : BrokerClient(connection) {
 
 }
 
-private class RatelimitProvider(private val burst: Int, private val duration: Duration) {
+private class KafkaRatelimitProvider(private val burst: Int, private val duration: Duration) {
 
-    private val limiters = ConcurrentHashMap<String, SuspendingRatelimit>()
+    private val limiters = ConcurrentHashMap<String, Ratelimit>()
 
-    fun getClientRatelimit(clientId: String): SuspendingRatelimit = limiters.computeIfAbsent(clientId) {
-        SuspendingRatelimit(burst, duration)
+    fun getClientRatelimit(clientId: String): Ratelimit = limiters.computeIfAbsent(clientId) {
+        Ratelimit(burst, duration)
     }
 
 }

vanilla/src/main/java/gg/beemo/vanilla/Vanilla.kt

@@ -5,6 +5,8 @@ import gg.beemo.latte.broker.rabbitmq.RabbitConnection
 import gg.beemo.latte.config.Configurator
 import gg.beemo.latte.logging.Log
 import gg.beemo.latte.logging.log
+import io.grpc.Server
+import io.grpc.ServerBuilder
 import kotlinx.coroutines.runBlocking
 import org.apache.logging.log4j.LogManager
 
@@ -22,19 +24,26 @@ object Vanilla {
         val brokerConnection = RabbitConnection(
             rabbitHosts = Config.RABBIT_HOST,
             serviceName = CommonConfig.BrokerServices.VANILLA,
-            instanceId  = "0", // There will only ever be one instance of vanilla
+            instanceId = "0", // There will only ever be one instance of vanilla
             useTls = Config.RABBIT_USE_TLS,
             username = Config.RABBIT_USERNAME,
             password = Config.RABBIT_PASSWORD,
         )
 
         log.debug("Initializing Kafka Ratelimit client")
-        val ratelimitClient = RatelimitClient(brokerConnection)
+        val ratelimitClient = KafkaRatelimitClient(brokerConnection)
+
+        val grpcServer: Server = ServerBuilder.forPort(Config.GRPC_PORT)
+            .addService(GrpcRatelimitService())
+            .build()
+            .start()
 
         Runtime.getRuntime().addShutdownHook(Thread({
             log.info("Destroying everything")
             ratelimitClient.destroy()
             brokerConnection.destroy()
+            grpcServer.shutdown()
+            grpcServer.awaitTermination()
             LogManager.shutdown(true, true)
         }, "Vanilla Shutdown Hook"))
 

vanilla/src/main/proto/ratelimit.proto

@@ -0,0 +1 @@
+../../../../proto/ratelimit.proto