Implement dispatching to custom remote functions

ballerina/annotation.bal

@@ -28,12 +28,14 @@
 #   If this is not set or is negative or zero, the default frame size, which is 65536 will be used
 # + auth - Listener authentication configurations
 # + validation - Enable/disable constraint validation
+# + dispatcherKey - The key which is going to be used for dispatching to custom remote functions.
 public type WSServiceConfig record {|
     string[] subProtocols = [];
     decimal idleTimeout = 0;
     int maxFrameSize = 65536;
     ListenerAuthConfig[] auth?;
     boolean validation = true;
+    string dispatcherKey?;
 |};
 
 # The annotation which is used to configure a WebSocket service.

ballerina/tests/custom_remote_functions.bal

@@ -0,0 +1,189 @@
+// Copyright (c) 2023 WSO2 LLC. (//www.wso2.org) All Rights Reserved.
+//
+// WSO2 Inc. 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
+//
+// //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;
+import ballerina/io;
+
+listener Listener customDispatchingLis = new(21401);
+
+@ServiceConfig {dispatcherKey: "event"}
+service on customDispatchingLis {
+    resource function get .() returns Service|Error {
+        return new SingleMessageService();
+    }
+}
+
+service class SingleMessageService {
+    *Service;
+    remote function onPing(Caller caller, json data) returns Error? {
+        io:println(data);
+        check caller->writeMessage({"event": "pong"});
+    }
+}
+
+@ServiceConfig {dispatcherKey: "event"}
+service /subscribe on customDispatchingLis {
+    resource function get .() returns Service|Error {
+        return new SubscribeMessageService();
+    }
+}
+
+service class SubscribeMessageService {
+    *Service;
+    remote function onSubscribe(Caller caller, json data) returns Error? {
+        io:println(data);
+        check caller->writeMessage({"type": "subscribe", "id":"1", "payload":{"query": "{ __schema { types { name } } }"}});
+    }
+
+    remote function onHeartbeat(Caller caller, json data) returns json {
+        return {"event": "heartbeat"};
+    }
+}
+
+@ServiceConfig {dispatcherKey: "event"}
+service /onMessage on customDispatchingLis {
+    resource function get .() returns Service|Error {
+        return new DefaultRemoteFunctionService();
+    }
+}
+
+service class DefaultRemoteFunctionService {
+    *Service;
+    remote function onMessage(Caller caller, json data) returns Error? {
+        check caller->writeMessage({"event": "heartbeat"});
+    }
+}
+
+@ServiceConfig {dispatcherKey: "event"}
+service /noRemoteMethod on customDispatchingLis {
+    resource function get .() returns Service|Error {
+        return new NoRemoteFunctionService();
+    }
+}
+
+service class NoRemoteFunctionService {
+    *Service;
+    remote function onMessages(Caller caller, json data) returns Error? {
+        check caller->writeMessage({"event": "onMessages"});
+    }
+}
+
+@ServiceConfig {dispatcherKey: "event"}
+service /dataBindingFailure on customDispatchingLis {
+ resource function get .() returns Service|Error {
+     return new DataBindingFailureService();
+ }
+}
+
+service class DataBindingFailureService {
+ *Service;
+ remote function onMessages(Caller caller, byte[] data) returns Error? {
+     check caller->writeMessage({"event": "onMessages"});
+ }
+}
+
+@ServiceConfig {dispatcherKey: "type"}
+service /underscore on customDispatchingLis {
+    resource function get .() returns Service|Error {
+     return new UnderscoreService();
+    }
+}
+
+service class UnderscoreService {
+    *Service;
+    remote function onPing(Caller caller, json data) returns Error? {
+     check caller->writeMessage({"event": "onMessages"});
+    }
+}
+
+@ServiceConfig {dispatcherKey: "type"}
+service /lotofunderscores on customDispatchingLis {
+    resource function get .() returns Service|Error {
+     return new LotOfUnderscoreService();
+    }
+}
+
+service class LotOfUnderscoreService {
+    *Service;
+    remote function onThisPingMessage(Caller caller, json data) returns Error? {
+     check caller->writeMessage({"event": "onMessages"});
+    }
+}
+
+@test:Config {}
+public function testPingMessage() returns Error? {
+    Client cl = check new("ws://localhost:21401");
+    check cl->writeMessage({"event": "ping"});
+    json resp = check cl->readMessage();
+    test:assertEquals(resp, {"event": "pong"});
+}
+
+@test:Config {}
+public function testSubscribeMessage() returns Error? {
+    Client cl = check new("ws://localhost:21401/subscribe");
+    check cl->writeMessage({"event": "subscribe", "pair": ["XBT/USD", "XBT/EUR"], "subscription": {"name": "ticker"}});
+    json resp = check cl->readMessage();
+    test:assertEquals(resp, {"type": "subscribe", "id":"1", "payload":{"query": "{ __schema { types { name } } }"}});
+    check cl->writeMessage({"event": "heartbeat"});
+    json resp2 = check cl->readMessage();
+    test:assertEquals(resp2, {"event": "heartbeat"});
+}
+
+@test:Config {}
+public function testDispatchingToDefaultRemoteMethod() returns Error? {
+    Client cl = check new("ws://localhost:21401/onMessage");
+    check cl->writeMessage({"event": "heartbeat"});
+    check cl->writeMessage({"event": "heartbeat"});
+    json resp2 = check cl->readMessage();
+    test:assertEquals(resp2, {"event": "heartbeat"});
+}
+
+@test:Config {}
+public function testDispatchingToNone() returns Error? {
+    Client cl = check new("ws://localhost:21401/noRemoteMethod");
+    check cl->writeMessage({"event": "heartbeat"});
+    check cl->writeMessage({"event": "Messages"});
+    json resp2 = check cl->readMessage();
+    test:assertEquals(resp2, {"event": "onMessages"});
+}
+
+@test:Config {}
+public function testDatabindingFailure() returns Error? {
+    Client cl = check new("ws://localhost:21401/dataBindingFailure");
+    check cl->writeMessage({"event": "Messages"});
+    json|Error resp2 = cl->readMessage();
+    if resp2 is Error {
+        test:assertTrue(resp2.message().startsWith("data binding failed:"));
+    } else {
+        test:assertFail("Expected a binding error");
+    }
+}
+
+@test:Config {}
+public function testUnderscore() returns Error? {
+    Client cl = check new("ws://localhost:21401/underscore");
+    check cl->writeMessage({"type": "_ping"});
+    json resp = check cl->readMessage();
+    test:assertEquals(resp, {"event": "onMessages"});
+}
+
+@test:Config {}
+public function testUnderscoresAndSpaces() returns Error? {
+    Client cl = check new("ws://localhost:21401/lotofunderscores");
+    check cl->writeMessage({"type": "this_ping message"});
+    json resp = check cl->readMessage();
+    test:assertEquals(resp, {"event": "onMessages"});
+}

ballerina/tests/data_binding_with_enum.bal

@@ -22,7 +22,7 @@ enum Point {
     C
 }
 
-service on new Listener(8000) {
+service on new Listener(21400) {
     resource function get .() returns Service|Error {
         return new EnumService();
     }
@@ -38,7 +38,7 @@ service class EnumService {
 
 @test:Config {}
 public function testEnum() returns Error? {
-    Client cl = check new("ws://localhost:8000");
+    Client cl = check new("ws://localhost:21400");
     check cl->writeMessage(A);
     Point p = check cl->readMessage();
     test:assertEquals(p, B);

changelog.md

@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+### Added
+- [Implement dispatching to custom remote functions](https://github.com/ballerina-platform/ballerina-standard-library/issues/3670)
+
 ### Fixed
 - [Fix data binding failure when a union type is used](https://github.com/ballerina-platform/ballerina-standard-library/issues/3909)
 - [Fix data binding failure when a enum type is used](https://github.com/ballerina-platform/ballerina-standard-library/issues/3707)

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

@@ -517,6 +517,14 @@ public void testReturningStreamTypeFromOnOpen() {
         Assert.assertEquals(diagnosticResult.errorCount(), 0);
     }
 
+    @Test
+    public void testCustomRemoteFunctions() {
+        Package currentPackage = loadPackage("sample_package_57");
+        PackageCompilation compilation = currentPackage.getCompilation();
+        DiagnosticResult diagnosticResult = compilation.diagnosticResult();
+        Assert.assertEquals(diagnosticResult.errorCount(), 0);
+    }
+
     private void assertDiagnostic(Diagnostic diagnostic, PluginConstants.CompilationErrors error) {
         Assert.assertEquals(diagnostic.diagnosticInfo().code(), error.getErrorCode());
         Assert.assertEquals(diagnostic.diagnosticInfo().messageFormat(),

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

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

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

@@ -0,0 +1,40 @@
+// Copyright (c) 2023 WSO2 LLC. (www.wso2.com) All Rights Reserved.
+//
+// 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
+//
+// //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;
+
+listener websocket:Listener localListener = new(8080);
+
+@websocket:ServiceConfig {dispatcherKey: "type"}
+service / on localListener {
+   resource function get .() returns websocket:Service|error {
+       return new WsService();
+   }
+}
+
+service class WsService {
+    *websocket:Service;
+
+    remote isolated function onOpen123(websocket:Caller caller) returns stream<string>|error {
+        string[] greets = ["Hi Sam", "Hey Sam", "GM Sam"];
+        return greets.toStream();
+    }
+
+    remote isolated function onMessagexxx(json data) returns stream<string, error?>|error {
+        string[] greets = ["Hi Sam", "Hey Sam", "GM Sam"];
+        return greets.toStream();
+    }
+}

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

@@ -41,6 +41,7 @@ public class PluginConstants {
     static final String RESOURCE_KEY_WORD = "resource";
     static final String HTTP_REQUEST = "http:Request";
     static final String LISTENER_IDENTIFIER = "Listener";
+    static final String DISPATCHER_ANNOTATION = "dispatcherKey";
 
     /**
      * Compilation Errors of WebSocket module.

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

@@ -135,7 +135,7 @@ private void filterRemoteFunctions(FunctionDefinitionNode functionDefinitionNode
                 Utils.validateOnBinaryMessageFunction(functionTypeSymbol, ctx, functionDefinitionNode);
                 break;
             default:
-                reportInvalidFunction(functionDefinitionNode);
+                // no need to validate
         }
     }
 

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

@@ -56,7 +56,7 @@ public class WebSocketUpgradeServiceValidator {
         this.modulePrefix = modulePrefix;
     }
 
-    void validate() {
+    void validate(boolean customDispatchingEnabled) {
         ServiceDeclarationNode serviceDeclarationNode = (ServiceDeclarationNode) ctx.node();
         boolean hasRemoteService = serviceDeclarationNode.members().stream().anyMatch(child ->
                 child.kind() == SyntaxKind.RESOURCE_ACCESSOR_DEFINITION);
@@ -93,6 +93,9 @@ void validate() {
         });
         if (resourceNode != null) {
             validateResourceReturnTypes(resourceNode);
+            if (customDispatchingEnabled) {
+                return;
+            }
             ReturnStatementNodeVisitor returnStatementNodeVisitor = new ReturnStatementNodeVisitor();
             resourceNode.accept(returnStatementNodeVisitor);
             for (ReturnStatementNode returnStatementNode : returnStatementNodeVisitor.getReturnStatementNodes()) {