Added safe, idempotent, isSampledToLocalTracing and method name parameters

Author: Mr3zee

compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcDeclarationScanner.kt

@@ -41,6 +41,7 @@ internal object RpcDeclarationScanner {
                     }
 
                     ServiceDeclaration.Method(
+                        ctx = ctx,
                         function = declaration,
                         arguments = ctx.versionSpecificApi.run {
                             declaration.valueParametersVS().memoryOptimizedMap { param ->

compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcIrContext.kt

@@ -64,6 +64,14 @@ internal class RpcIrContext(
         getRpcIrClassSymbol("RpcServiceDescriptor", "descriptor")
     }
 
+    val grpcAnnotation by lazy {
+        getIrClassSymbol("kotlinx.rpc.grpc.annotations", "Grpc")
+    }
+
+    val grpcMethodAnnotation by lazy {
+        grpcAnnotation.subClass("Method")
+    }
+
     val grpcServiceDescriptor by lazy {
         getIrClassSymbol("kotlinx.rpc.grpc.descriptor", "GrpcServiceDescriptor")
     }

compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/RpcStubGenerator.kt

@@ -7,7 +7,6 @@ package kotlinx.rpc.codegen.extension
 import kotlinx.rpc.codegen.VersionSpecificApi
 import kotlinx.rpc.codegen.common.RpcClassId
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
-import org.jetbrains.kotlin.backend.common.lower.irThrow
 import org.jetbrains.kotlin.backend.jvm.functionByName
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
@@ -1076,7 +1075,7 @@ internal class RpcStubGenerator(
             keyType = ctx.irBuiltIns.stringType,
             valueType = ctx.grpcPlatformMethodDescriptor.starProjectedType,
             declaration.methods.memoryOptimizedMap { callable ->
-                stringConst(callable.name) to irMethodDescriptor(callable, resolver)
+                stringConst(callable.grpcName) to irMethodDescriptor(callable, resolver)
             },
         )
     }
@@ -1130,15 +1129,14 @@ internal class RpcStubGenerator(
      * // In scope: resolver: MessageCodecResolver
      *
      * methodDescriptor<<request-type>, <response-type>>(
-     *     fullMethodName = "${descriptor.serviceFqName}/${callable.name}",
+     *     fullMethodName = "${descriptor.serviceFqName}/${<from Grpc.Method annotation> ?: callable.name}",
      *     requestCodec = <request-codec>,
      *     responseCodec = <response-codec>,
      *     type = MethodType.<method-type>,
      *     schemaDescriptor = null, // null for now
-     *     // todo understand these values
-     *     idempotent = true,
-     *     safe = true,
-     *     sampledToLocalTracing = true,
+     *     idempotent = <from Grpc.Method annotation>,
+     *     safe = <from Grpc.Method annotation>,
+     *     sampledToLocalTracing = <from Grpc.Method annotation>,
      * )
      * ```
      *
@@ -1183,9 +1181,12 @@ internal class RpcStubGenerator(
                     +responseType
                 }
 
+                val grpcMethodAnnotation = callable.function
+                    .getAnnotation(ctx.grpcMethodAnnotation.owner.kotlinFqName)
+
                 values {
                     // fullMethodName
-                    +stringConst("${declaration.serviceFqName}/${callable.name}")
+                    +stringConst("${declaration.serviceFqName}/${callable.grpcName}")
 
                     // requestCodec
                     +irCodec(requestType, resolver)
@@ -1220,15 +1221,26 @@ internal class RpcStubGenerator(
                     // schemaDescriptor
                     +nullConst(ctx.anyNullable)
 
-                    // todo figure out these
+                    val idempotentArgument = grpcMethodAnnotation
+                        ?.getValueArgument(Name.identifier("idempotent"))
+                    val idempotent = (idempotentArgument as? IrConst)?.value as? Boolean ?: false
+
                     // idempotent
-                    +booleanConst(true)
+                    +booleanConst(idempotent)
+
+                    val safeArgument = grpcMethodAnnotation
+                        ?.getValueArgument(Name.identifier("safe"))
+                    val safe = (safeArgument as? IrConst)?.value as? Boolean ?: false
 
                     // safe
-                    +booleanConst(true)
+                    +booleanConst(safe)
+
+                    val sampledToLocalTracingArgument = grpcMethodAnnotation
+                        ?.getValueArgument(Name.identifier("sampledToLocalTracing"))
+                    val sampledToLocalTracing = (sampledToLocalTracingArgument as? IrConst)?.value as? Boolean ?: true
 
                     // sampledToLocalTracing
-                    +booleanConst(true)
+                    +booleanConst(sampledToLocalTracing)
                 }
             }
         }

compiler-plugin/compiler-plugin-backend/src/main/kotlin/kotlinx/rpc/codegen/extension/ServiceDeclaration.kt

@@ -8,12 +8,16 @@ import kotlinx.rpc.codegen.common.RpcClassId
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.ir.util.getAnnotation
+import org.jetbrains.kotlin.ir.util.getValueArgument
 import org.jetbrains.kotlin.ir.util.hasAnnotation
 import org.jetbrains.kotlin.ir.util.kotlinFqName
+import org.jetbrains.kotlin.name.Name
 
-class ServiceDeclaration(
+internal class ServiceDeclaration(
     val service: IrClass,
     val stubClass: IrClass,
     val methods: List<Method>,
@@ -33,10 +37,22 @@ class ServiceDeclaration(
     }
 
     class Method(
+        val ctx: RpcIrContext,
         val function: IrSimpleFunction,
         val arguments: List<Argument>,
     ) : Callable {
         override val name: String = function.name.asString()
+        val grpcName by lazy {
+            val grpcMethodAnnotation = function.getAnnotation(
+                ctx.grpcMethodAnnotation.owner.kotlinFqName
+            )
+
+            val nameArgument = grpcMethodAnnotation?.getValueArgument(Name.identifier("name"))
+
+            ((nameArgument as? IrConst)?.value as? String)
+                ?.takeIf { it.isNotBlank() }
+                ?: name
+        }
 
         class Argument(
             val value: IrValueParameter,

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/annotations/Grpc.kt

@@ -11,4 +11,12 @@ import kotlinx.rpc.annotations.Rpc
  */
 @Target(AnnotationTarget.CLASS, AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.TYPE_PARAMETER)
 @Rpc
-public annotation class Grpc(val protoPackage: String = "")
+public annotation class Grpc(val protoPackage: String = "") {
+    @Target(AnnotationTarget.FUNCTION)
+    public annotation class Method(
+        val name: String = "",
+        val safe: Boolean = false,
+        val idempotent: Boolean = false,
+        val sampledToLocalTracing: Boolean = true,
+    )
+}

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/GrpcAnnotationsTest.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2023-2025 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.rpc.grpc.test
+
+import kotlinx.rpc.descriptor.serviceDescriptorOf
+import kotlinx.rpc.grpc.annotations.Grpc
+import kotlinx.rpc.grpc.codec.EmptyMessageCodecResolver
+import kotlinx.rpc.grpc.codec.MessageCodec
+import kotlinx.rpc.grpc.codec.MessageCodecResolver
+import kotlinx.rpc.grpc.descriptor.GrpcServiceDescriptor
+import kotlinx.rpc.protobuf.input.stream.InputStream
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertFalse
+import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
+
+@Grpc("grpc.test")
+interface GrpcAnnotationsService {
+    @Suppress("unused")
+    @Grpc.Method(name = "Empty", safe = true, idempotent = true, sampledToLocalTracing = false)
+    suspend fun empty()
+}
+
+class GrpcAnnotationsTest {
+    @Test
+    fun nullCodec() {
+        assertFailsWith<IllegalArgumentException> {
+            val descriptor = serviceDescriptorOf<GrpcAnnotationsService>()
+                    as GrpcServiceDescriptor<GrpcAnnotationsService>
+
+            descriptor.delegate(EmptyMessageCodecResolver)
+        }
+    }
+
+    @Test
+    fun methodAnnotations() {
+        val descriptor = serviceDescriptorOf<GrpcAnnotationsService>()
+                as GrpcServiceDescriptor<GrpcAnnotationsService>
+        val methodDescriptor = descriptor
+            .delegate(unitCodec)
+            .getMethodDescriptor("Empty")
+
+        assertNotNull(methodDescriptor)
+
+        assertEquals("grpc.test.GrpcAnnotationsService/Empty", methodDescriptor.getFullMethodName())
+        assertTrue(methodDescriptor.isSafe())
+        assertTrue(methodDescriptor.isIdempotent())
+        assertFalse(methodDescriptor.isSampledToLocalTracing())
+    }
+}
+
+private val unitCodec = MessageCodecResolver {
+    object : MessageCodec<Unit> {
+        override fun encode(value: Unit): InputStream {
+            TODO("Not yet implemented")
+        }
+
+        override fun decode(stream: InputStream) {
+            TODO("Not yet implemented")
+        }
+    }
+}
+