Add meta annotation support in Kotlin annotation processor (#431)

sdk-api-kotlin-gen/src/main/kotlin/dev/restate/sdk/kotlin/gen/KElementConverter.kt

@@ -66,47 +66,17 @@ class KElementConverter(
       logger.error("The annotated class is private", classDeclaration)
     }
 
-    // Figure out service type annotations
-    val serviceAnnotation =
-        classDeclaration
-            .getAnnotationsByType(dev.restate.sdk.annotation.Service::class)
-            .firstOrNull()
-    val virtualObjectAnnotation =
-        classDeclaration
-            .getAnnotationsByType(dev.restate.sdk.annotation.VirtualObject::class)
-            .firstOrNull()
-    val workflowAnnotation =
-        classDeclaration
-            .getAnnotationsByType(dev.restate.sdk.annotation.Workflow::class)
-            .firstOrNull()
-    val isAnnotatedWithService = serviceAnnotation != null
-    val isAnnotatedWithVirtualObject = virtualObjectAnnotation != null
-    val isAnnotatedWithWorkflow = workflowAnnotation != null
-
-    // Check there's exactly one annotation
-    if (!(isAnnotatedWithService xor isAnnotatedWithVirtualObject xor isAnnotatedWithWorkflow)) {
-      logger.error(
-          "The type can be annotated only with one annotation between @VirtualObject and @Service and @Workflow",
-          classDeclaration)
-    }
-
-    data.withServiceType(
-        if (isAnnotatedWithService) ServiceType.SERVICE
-        else if (isAnnotatedWithWorkflow) ServiceType.WORKFLOW else ServiceType.VIRTUAL_OBJECT)
-
     // Infer names
     val targetPkg = classDeclaration.packageName.asString()
     val targetFqcn = classDeclaration.qualifiedName!!.asString()
-    var serviceName =
-        if (isAnnotatedWithService) serviceAnnotation!!.name
-        else if (isAnnotatedWithWorkflow) workflowAnnotation!!.name
-        else virtualObjectAnnotation!!.name
-    if (serviceName.isEmpty()) {
+    data.withTargetPkg(targetPkg).withTargetFqcn(targetFqcn)
+
+    if (data.serviceName.isNullOrEmpty()) {
       // Use Simple class name
       // With this logic we make sure we flatten subclasses names
-      serviceName = targetFqcn.substring(targetPkg.length).replace(Pattern.quote(".").toRegex(), "")
+      data.withServiceName(
+          targetFqcn.substring(targetPkg.length).replace(Pattern.quote(".").toRegex(), ""))
     }
-    data.withTargetPkg(targetPkg).withTargetFqcn(targetFqcn).withServiceName(serviceName)
 
     // Compute handlers
     classDeclaration

sdk-api-kotlin-gen/src/main/kotlin/dev/restate/sdk/kotlin/gen/MetaRestateAnnotation.kt

@@ -0,0 +1,25 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.kotlin.gen
+
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSName
+import dev.restate.sdk.common.ServiceType
+
+internal data class MetaRestateAnnotation(
+    val annotationName: KSName,
+    val serviceType: ServiceType
+) {
+  fun resolveName(annotated: KSAnnotated): String? =
+      annotated.annotations
+          .find { it.annotationType.resolve().declaration.qualifiedName == annotationName }
+          ?.arguments
+          ?.firstOrNull { it -> it.name?.getShortName() == "name" }
+          ?.value as String?
+}

sdk-api-kotlin-gen/src/main/kotlin/dev/restate/sdk/kotlin/gen/ServiceProcessor.kt

@@ -11,8 +11,10 @@ package dev.restate.sdk.kotlin.gen
 import com.github.jknack.handlebars.io.ClassPathTemplateLoader
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.containingFile
+import com.google.devtools.ksp.getClassDeclarationByName
 import com.google.devtools.ksp.getKotlinClassByName
 import com.google.devtools.ksp.processing.*
+import com.google.devtools.ksp.symbol.ClassKind
 import com.google.devtools.ksp.symbol.KSAnnotated
 import com.google.devtools.ksp.symbol.KSClassDeclaration
 import com.google.devtools.ksp.symbol.Origin
@@ -68,35 +70,24 @@ class ServiceProcessor(private val logger: KSPLogger, private val codeGenerator:
             resolver.builtIns,
             resolver.getKotlinClassByName(ByteArray::class.qualifiedName!!)!!.asType(listOf()))
 
-    val resolved =
-        resolver
-            .getSymbolsWithAnnotation(dev.restate.sdk.annotation.Service::class.qualifiedName!!)
-            .toSet() +
-            resolver
-                .getSymbolsWithAnnotation(
-                    dev.restate.sdk.annotation.VirtualObject::class.qualifiedName!!)
-                .toSet() +
-            resolver
-                .getSymbolsWithAnnotation(
-                    dev.restate.sdk.annotation.Workflow::class.qualifiedName!!)
-                // Workflow annotation can be set on functions too
-                .filter { ksAnnotated -> ksAnnotated is KSClassDeclaration }
-                .toSet()
+    val discovered = discoverRestateAnnotatedOrMetaAnnotatedServices(resolver)
 
     val services =
-        resolved
-            .filter { it.containingFile!!.origin == Origin.KOTLIN }
+        discovered
             .map {
               val serviceBuilder = Service.builder()
-              converter.visitAnnotated(it, serviceBuilder)
+              serviceBuilder.withServiceType(it.first.serviceType)
+              serviceBuilder.withServiceName(it.first.resolveName(it.second))
+
+              converter.visitAnnotated(it.second, serviceBuilder)
 
               var serviceModel: Service? = null
               try {
                 serviceModel = serviceBuilder.validateAndBuild()
               } catch (e: Exception) {
-                logger.error("Unable to build service: $e", it)
+                logger.error("Unable to build service: $e", it.second)
               }
-              (it to serviceModel!!)
+              (it.second to serviceModel!!)
             }
             .toList()
 
@@ -127,6 +118,80 @@ class ServiceProcessor(private val logger: KSPLogger, private val codeGenerator:
     return emptyList()
   }
 
+  private fun discoverRestateAnnotatedOrMetaAnnotatedServices(
+      resolver: Resolver
+  ): Set<Pair<MetaRestateAnnotation, KSAnnotated>> {
+    val discoveredAnnotatedElements = mutableSetOf<Pair<MetaRestateAnnotation, KSAnnotated>>()
+
+    val metaAnnotationsToProcess =
+        mutableListOf(
+            MetaRestateAnnotation(
+                resolver
+                    .getClassDeclarationByName<dev.restate.sdk.annotation.Service>()!!
+                    .qualifiedName!!,
+                ServiceType.SERVICE),
+            MetaRestateAnnotation(
+                resolver
+                    .getClassDeclarationByName<dev.restate.sdk.annotation.VirtualObject>()!!
+                    .qualifiedName!!,
+                ServiceType.VIRTUAL_OBJECT),
+            MetaRestateAnnotation(
+                resolver
+                    .getClassDeclarationByName<dev.restate.sdk.annotation.Workflow>()!!
+                    .qualifiedName!!,
+                ServiceType.WORKFLOW))
+    val discoveredAnnotations = mutableSetOf<String>()
+
+    var metaAnnotation = metaAnnotationsToProcess.removeFirstOrNull()
+    while (metaAnnotation != null) {
+      if (!discoveredAnnotations.add(metaAnnotation.annotationName.asString())) {
+        // We alredy discovered it, skip
+        continue
+      }
+      for (annotatedElement in
+          resolver.getSymbolsWithAnnotation(metaAnnotation.annotationName.asString())) {
+        if (annotatedElement !is KSClassDeclaration) {
+          continue
+        }
+        when (annotatedElement.classKind) {
+          ClassKind.INTERFACE,
+          ClassKind.CLASS -> {
+            if (annotatedElement.containingFile!!.origin != Origin.KOTLIN) {
+              // Skip if it's not kotlin
+              continue
+            }
+            discoveredAnnotatedElements.add(metaAnnotation to annotatedElement)
+          }
+          ClassKind.ANNOTATION_CLASS -> {
+            metaAnnotationsToProcess.add(
+                MetaRestateAnnotation(annotatedElement.qualifiedName!!, metaAnnotation.serviceType))
+          }
+          else ->
+              logger.error(
+                  "The ServiceProcessor supports only interfaces or classes declarations",
+                  annotatedElement)
+        }
+      }
+      metaAnnotation = metaAnnotationsToProcess.removeFirstOrNull()
+    }
+
+    val knownAnnotations = discoveredAnnotations.toSet()
+
+    // Check annotated elements are annotated with only one of the given annotations.
+    discoveredAnnotatedElements.forEach { it ->
+      val forbiddenAnnotations = knownAnnotations - setOf(it.first.annotationName.asString())
+      val elementAnnotations =
+          it.second.annotations
+              .mapNotNull { it.annotationType.resolve().declaration.qualifiedName?.asString() }
+              .toSet()
+      if (forbiddenAnnotations.intersect(elementAnnotations).isNotEmpty()) {
+        logger.error("The type is annotated with more than one Restate annotation", it.second)
+      }
+    }
+
+    return discoveredAnnotatedElements.toSet()
+  }
+
   private fun generateMetaINF(services: List<Pair<KSAnnotated, Service>>) {
     val resourceFile = "META-INF/services/${ServiceDefinitionFactory::class.java.canonicalName}"
     val dependencies =

sdk-api-kotlin-gen/src/test/kotlin/dev/restate/sdk/kotlin/CodegenTest.kt

@@ -202,6 +202,15 @@ class CodegenTest : TestDefinitions.TestSuite {
     }
   }
 
+  // Just needs to compile
+  @MyMetaServiceAnnotation(name = "MetaAnnotatedGreeter")
+  class MetaAnnotatedGreeter {
+    @Handler
+    suspend fun greet(context: Context, request: String): String {
+      return MetaAnnotatedGreeterClient.fromContext(context).greet(request).await()
+    }
+  }
+
   override fun definitions(): Stream<TestDefinition> {
     return Stream.of(
         testInvocation({ ServiceGreeter() }, "greet")

sdk-api-kotlin-gen/src/test/kotlin/dev/restate/sdk/kotlin/MyMetaServiceAnnotation.kt

@@ -0,0 +1,13 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.kotlin
+
+import dev.restate.sdk.annotation.Service
+
+@Service annotation class MyMetaServiceAnnotation(val name: String = "")