Add support for custom remote function mapping via Annotation

ballerina/CompilerPlugin.toml

@@ -4,3 +4,9 @@ class = "io.ballerina.stdlib.websocket.plugin.WebSocketCompilerPlugin"
 
 [[dependency]]
 path = "../compiler-plugin/build/libs/websocket-compiler-plugin-2.15.0-SNAPSHOT.jar"
+
+[[dependency]]
+path = "../native/build/libs/websocket-native-2.15.0-SNAPSHOT.jar"
+
+[[dependency]]
+path = "./lib/http-native-2.14.0.jar"

ballerina/annotation.bal

@@ -14,10 +14,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-///////////////////////////
-/// Service Annotations ///
-///////////////////////////
-
 # Configurations for a WebSocket service.
 #
 # + subProtocols - Negotiable sub protocol by the service
@@ -48,3 +44,13 @@ public type WSServiceConfig record {|
 
 # The annotation which is used to configure a WebSocket service.
 public annotation WSServiceConfig ServiceConfig on service;
+
+# Configurations used to define dispatching rules for remote functions.
+#
+# + value - The value which is going to be used for dispatching to custom remote functions.
+public type WsDispatcherMapping record {|
+    string value;
+|};
+
+# The annotation which is used to configure the dispatching rules for WebSocket remote functions.
+public const annotation WsDispatcherMapping DispatcherMapping on function;

ballerina/build.gradle

@@ -122,6 +122,7 @@ task updateTomlFiles {
         ballerinaTomlFile.text = newConfig
 
         def newPluginConfig = compilerPluginTomlFilePlaceHolder.text.replace("@project.version@", project.version)
+        newPluginConfig = newPluginConfig.replace("@stdlib.httpnative.version@", stdlibDependentHttpNativeVersion)
         compilerPluginTomlFile.text = newPluginConfig
     }
 }

ballerina/tests/test_dispatcher_mapping_annotation.bal

@@ -0,0 +1,66 @@
+//  Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+//
+//  WSO2 LLC. licenses this file to you under the Apache License,
+//  Version 2.0 (the "License"); you may not use this file except
+//  in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing,
+//  software distributed under the License is distributed on an
+//  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+//  KIND, either express or implied. See the License for the
+//  specific language governing permissions and limitations
+//  under the License.
+
+import ballerina/test;
+
+type Subscribe record {|
+    string event = "subscribe";
+    string data;
+|};
+
+@ServiceConfig {
+    dispatcherKey: "event"
+}
+service / on new Listener(22103) {
+    resource function get .() returns Service|UpgradeError {
+        return new WsService22103();
+    }
+}
+
+service class WsService22103 {
+    *Service;
+
+    @DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeMessage(Subscribe message) returns string {
+        return "onSubscribeMessage";
+    }
+
+    remote function onSubscribeMessageError(Caller caller, error message) returns error? {
+        check caller->writeMessage("onSubscribeMessageError");
+    }
+}
+
+@test:Config {
+    groups: ["dispatcherMappingAnnotation"]
+}
+public function testDispatcherMappingAnnotation() returns error? {
+    Client wsClient = check new ("ws://localhost:22103/");
+    check wsClient->writeMessage({event: "subscribe", data: "test"});
+    string res = check wsClient->readMessage();
+    test:assertEquals(res, "onSubscribeMessage");
+}
+
+@test:Config {
+    groups: ["dispatcherMappingAnnotation"]
+}
+public function testDispatcherMappingAnnotationWithCustomOnError() returns error? {
+    Client wsClient = check new ("ws://localhost:22103/");
+    check wsClient->writeMessage({event: "subscribe", invalidField: "test"});
+    string res = check wsClient->readMessage();
+    test:assertEquals(res, "onSubscribeMessageError");
+}

build-config/resources/CompilerPlugin.toml

@@ -4,3 +4,9 @@ class = "io.ballerina.stdlib.websocket.plugin.WebSocketCompilerPlugin"
 
 [[dependency]]
 path = "../compiler-plugin/build/libs/websocket-compiler-plugin-@project.version@.jar"
+
+[[dependency]]
+path = "../native/build/libs/websocket-native-@project.version@.jar"
+
+[[dependency]]
+path = "./lib/http-native-@stdlib.httpnative.version@.jar"

changelog.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Added
 
+- [Support Custom Remote Function Mapping via Annotation](https://github.com/ballerina-platform/ballerina-library/issues/7733)
 - [Introduce Service Config Annotation for connectionClosureTimeout in Websocket module](https://github.com/ballerina-platform/ballerina-library/issues/7697)
 - [Implement websocket close frame support](https://github.com/ballerina-platform/ballerina-library/issues/7578)
 

compiler-plugin-tests/src/test/java/io/ballerina/stdlib/websocket/compiler/WebSocketServiceValidationTest.java

@@ -572,6 +572,20 @@ public void testRemoteFunctionWithStreamAndCloseFrameReturnTypes() {
         Assert.assertEquals(diagnosticResult.errorCount(), 0);
     }
 
+    @Test
+    public void testDispatcherMappingAnnotation() {
+        Package currentPackage = loadPackage("sample_package_63");
+        PackageCompilation compilation = currentPackage.getCompilation();
+        DiagnosticResult diagnosticResult = compilation.diagnosticResult();
+        Assert.assertEquals(diagnosticResult.errorCount(), 3);
+        Diagnostic firstDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[0];
+        assertDiagnostic(firstDiagnostic, PluginConstants.CompilationErrors.RE_DECLARED_REMOTE_FUNCTIONS);
+        Diagnostic secondDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[1];
+        assertDiagnostic(secondDiagnostic, PluginConstants.CompilationErrors.DUPLICATED_DISPATCHER_MAPPING_VALUE);
+        Diagnostic thirdDiagnostic = (Diagnostic) diagnosticResult.errors().toArray()[2];
+        assertDiagnostic(thirdDiagnostic, PluginConstants.CompilationErrors.INVALID_FUNCTION_ANNOTATION);
+    }
+
     @Test
     public void testConnectionClosureTimeoutInTheServiceConfig() {
         Package currentPackage = loadPackage("sample_package_64");

compiler-plugin-tests/src/test/resources/ballerina_sources/sample_package_63/Ballerina.toml

@@ -0,0 +1,4 @@
+[package]
+org = "websocket_test"
+name = "sample_63"
+version = "0.1.0"

compiler-plugin-tests/src/test/resources/ballerina_sources/sample_package_63/server.bal

@@ -0,0 +1,60 @@
+//  Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+//
+//  WSO2 LLC. licenses this file to you under the Apache License,
+//  Version 2.0 (the "License"); you may not use this file except
+//  in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing,
+//  software distributed under the License is distributed on an
+//  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+//  KIND, either express or implied. See the License for the
+//  specific language governing permissions and limitations
+//  under the License.
+
+import ballerina/websocket as ws;
+
+type Subscribe record {|
+    string event = "subscribe";
+    string data;
+|};
+
+@ws:ServiceConfig {
+    dispatcherKey: "event"
+}
+service / on new ws:Listener(9090) {
+    resource function get .() returns ws:Service|ws:UpgradeError {
+        return new WsService();
+    }
+}
+
+service class WsService {
+    *ws:Service;
+
+    remote function onSubscribe(Subscribe message) returns string {
+        return "onSubscribe";
+    }
+
+    @ws:DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeMessage(Subscribe message) returns string {
+        return "onSubscribeMessage";
+    }
+
+    @ws:DispatcherMapping {
+        value: "subscribe"
+    }
+    remote function onSubscribeText(Subscribe message) returns string {
+        return "onSubscribeText";
+    }
+
+    @ws:DispatcherMapping {
+        value: "ping"
+    }
+    remote function onPing(Subscribe message) returns string {
+        return "onPing";
+    }
+}

compiler-plugin/src/main/java/io/ballerina/stdlib/websocket/plugin/PluginConstants.java

@@ -79,6 +79,12 @@ public enum CompilationErrors {
                 "WEBSOCKET_215"),
         INVALID_REMOTE_FUNCTIONS("Cannot have `{0}` with `onMessage` remote function",
                 "WEBSOCKET_216"),
+        RE_DECLARED_REMOTE_FUNCTIONS("Cannot have `{0}` because the message type `{1}` is already " +
+                "associated with `{2}` remote function", "WEBSOCKET_217"),
+        DUPLICATED_DISPATCHER_MAPPING_VALUE("DispatcherMapping annotation value `{0}` is already " +
+                "exists", "WEBSOCKET_218"),
+        INVALID_FUNCTION_ANNOTATION("Invalid annotation provided for `{0}` remote function. "
+                + "This annotation can only be used with the custom dispatcher functions", "WEBSOCKET_219"),
         INVALID_RESOURCE_ERROR("There should be only one `get` resource for the service",
                 "WEBSOCKET_101"),
         MORE_THAN_ONE_RESOURCE_PARAM_ERROR("There should be only http:Request as a parameter",

compiler-plugin/src/main/java/io/ballerina/stdlib/websocket/plugin/WebSocketServiceValidator.java

@@ -20,30 +20,85 @@
 import io.ballerina.compiler.api.symbols.FunctionTypeSymbol;
 import io.ballerina.compiler.api.symbols.MethodSymbol;
 import io.ballerina.compiler.api.symbols.Qualifier;
+import io.ballerina.compiler.api.symbols.Symbol;
+import io.ballerina.compiler.syntax.tree.AnnotationNode;
 import io.ballerina.compiler.syntax.tree.ClassDefinitionNode;
+import io.ballerina.compiler.syntax.tree.ExpressionNode;
 import io.ballerina.compiler.syntax.tree.FunctionDefinitionNode;
+import io.ballerina.compiler.syntax.tree.MappingConstructorExpressionNode;
+import io.ballerina.compiler.syntax.tree.Node;
 import io.ballerina.compiler.syntax.tree.NodeList;
+import io.ballerina.compiler.syntax.tree.SpecificFieldNode;
 import io.ballerina.compiler.syntax.tree.SyntaxKind;
 import io.ballerina.compiler.syntax.tree.Token;
 import io.ballerina.projects.plugins.SyntaxNodeAnalysisContext;
+import io.ballerina.stdlib.websocket.WebSocketConstants;
 import io.ballerina.tools.diagnostics.DiagnosticFactory;
 import io.ballerina.tools.diagnostics.DiagnosticInfo;
 import io.ballerina.tools.diagnostics.DiagnosticSeverity;
 
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
+import static io.ballerina.stdlib.websocket.WebSocketConstants.ANNOTATION_ATTR_DISPATCHER_VALUE;
+import static io.ballerina.stdlib.websocket.WebSocketResourceDispatcher.createCustomRemoteFunction;
+import static io.ballerina.stdlib.websocket.plugin.PluginConstants.CompilationErrors.DUPLICATED_DISPATCHER_MAPPING_VALUE;
+import static io.ballerina.stdlib.websocket.plugin.PluginConstants.CompilationErrors.INVALID_FUNCTION_ANNOTATION;
+import static io.ballerina.stdlib.websocket.plugin.PluginConstants.CompilationErrors.RE_DECLARED_REMOTE_FUNCTIONS;
+
 /**
  * A class for validating websocket service.
  */
 public class WebSocketServiceValidator {
     public static final String GENERIC_FUNCTION = "generic function";
+    private final Set<String> specialRemoteMethods = Set.of(PluginConstants.ON_OPEN, PluginConstants.ON_CLOSE,
+            PluginConstants.ON_ERROR, PluginConstants.ON_IDLE_TIMEOUT, PluginConstants.ON_TEXT_MESSAGE,
+            PluginConstants.ON_BINARY_MESSAGE, PluginConstants.ON_MESSAGE, PluginConstants.ON_PING_MESSAGE,
+            PluginConstants.ON_PONG_MESSAGE);
     private SyntaxNodeAnalysisContext ctx;
 
     WebSocketServiceValidator(SyntaxNodeAnalysisContext syntaxNodeAnalysisContext) {
         this.ctx = syntaxNodeAnalysisContext;
     }
 
+    private static Optional<String> getDispatcherMappingAnnotatedFunctionName(FunctionDefinitionNode node,
+                                                                              SyntaxNodeAnalysisContext ctx) {
+        if (node.metadata().isEmpty()) {
+            return Optional.empty();
+        }
+        for (AnnotationNode annotationNode : node.metadata().get().annotations()) {
+            Optional<Symbol> annotationType = ctx.semanticModel().symbol(annotationNode);
+            if (annotationType.isEmpty()) {
+                continue;
+            }
+            if (!annotationType.get().getModule().flatMap(Symbol::getName)
+                    .orElse("").equals(WebSocketConstants.PACKAGE_WEBSOCKET) ||
+                    !annotationType.get().getName().orElse("")
+                            .equals(WebSocketConstants.WEBSOCKET_DISPATCHER_MAPPING_ANNOTATION)) {
+                continue;
+            }
+            if (annotationNode.annotValue().isEmpty()) {
+                return Optional.empty();
+            }
+            MappingConstructorExpressionNode annotationValue = annotationNode.annotValue().get();
+            for (Node field : annotationValue.fields()) {
+                if (!field.kind().equals(SyntaxKind.SPECIFIC_FIELD)) {
+                    continue;
+                }
+                String fieldName = ((SpecificFieldNode) field).fieldName().toString().strip();
+                Optional<ExpressionNode> filedValue = ((SpecificFieldNode) field).valueExpr();
+                if (!fieldName.equals(ANNOTATION_ATTR_DISPATCHER_VALUE) || filedValue.isEmpty()) {
+                    continue;
+                }
+                return Optional.of(filedValue.get().toString().replaceAll("\"", "").strip());
+            }
+        }
+        return Optional.empty();
+    }
+
     public void validate() {
         ClassDefinitionNode classDefNode = (ClassDefinitionNode) ctx.node();
         Map<String, Boolean> functionSet = classDefNode.members().stream().filter(child ->
@@ -70,7 +125,7 @@ public void validate() {
                     classDefNode.location(), PluginConstants.ON_TEXT_MESSAGE);
         }
         if (functionSet.containsKey(PluginConstants.ON_MESSAGE) &&
-                        functionSet.containsKey(PluginConstants.ON_BINARY_MESSAGE)) {
+                functionSet.containsKey(PluginConstants.ON_BINARY_MESSAGE)) {
             Utils.reportDiagnostics(ctx, PluginConstants.CompilationErrors.INVALID_REMOTE_FUNCTIONS,
                     classDefNode.location(), PluginConstants.ON_BINARY_MESSAGE);
         }
@@ -105,6 +160,45 @@ public void validate() {
                 !functionSet.containsKey(PluginConstants.ON_BINARY_MESSAGE)) {
             reportDiagnostic(classDefNode, PluginConstants.CompilationErrors.ON_MESSAGE_GENERATION_HINT);
         }
+        validateDispatcherMappingAnnotations(classDefNode, functionSet);
+    }
+
+    private void validateDispatcherMappingAnnotations(ClassDefinitionNode classDefNode,
+                                                      Map<String, Boolean> functionSet) {
+        Set<String> seenAnnotationValues = new HashSet<>();
+        for (Node node : classDefNode.members()) {
+            if (!node.kind().equals(SyntaxKind.OBJECT_METHOD_DEFINITION)) {
+                continue;
+            }
+            FunctionDefinitionNode funcDefinitionNode = (FunctionDefinitionNode) node;
+            if (funcDefinitionNode.qualifierList().stream()
+                    .noneMatch(token -> token.text().equals(Qualifier.REMOTE.getValue()))) {
+                continue;
+            }
+            Optional<String> funcName = ctx.semanticModel().symbol(funcDefinitionNode).flatMap(Symbol::getName);
+            Optional<String> annoDispatchingValue = getDispatcherMappingAnnotatedFunctionName(funcDefinitionNode, ctx);
+            if (funcName.isEmpty() || annoDispatchingValue.isEmpty()) {
+                continue;
+            }
+            if (seenAnnotationValues.contains(annoDispatchingValue.get())) {
+                Utils.reportDiagnostics(ctx, DUPLICATED_DISPATCHER_MAPPING_VALUE,
+                        funcDefinitionNode.location(), annoDispatchingValue.get());
+                continue;
+            }
+            seenAnnotationValues.add(annoDispatchingValue.get());
+            String customRemoteFunctionName = createCustomRemoteFunction(annoDispatchingValue.get());
+            if (this.specialRemoteMethods.contains(funcName.get())) {
+                Utils.reportDiagnostics(ctx, INVALID_FUNCTION_ANNOTATION, funcDefinitionNode.location(),
+                        funcName.get());
+                continue;
+            }
+            if (functionSet.containsKey(customRemoteFunctionName) &&
+                    !customRemoteFunctionName.equals(funcName.get()) &&
+                    !this.specialRemoteMethods.contains(customRemoteFunctionName)) {
+                Utils.reportDiagnostics(ctx, RE_DECLARED_REMOTE_FUNCTIONS, classDefNode.location(),
+                        customRemoteFunctionName, annoDispatchingValue.get(), funcName.get());
+            }
+        }
     }
 
     private void filterRemoteFunctions(FunctionDefinitionNode functionDefinitionNode) {