feat(server): disable UDP service by default

apache · Dec 20, 2022 · 8a9f5f2 · 8a9f5f2
1 parent 17bd733
commit 8a9f5f2
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 14 deletions.
diff --git a/src/runtime/nativerun.cpp b/src/runtime/nativerun.cpp
@@ -34,10 +34,13 @@
  */
 
 #include "runtime/nativerun.h"
+#include "utils/flags.h"
 
 namespace dsn {
 namespace tools {
 
+DSN_DECLARE_bool(enable_udp);
+
 void nativerun::install(service_spec &spec)
 {
     if (spec.env_factory_name == "")
@@ -59,19 +62,21 @@ void nativerun::install(service_spec &spec)
         cs2.message_buffer_block_size = 1024 * 64;
         spec.network_default_server_cfs[cs2] = cs2;
     }
-    {
-        network_client_config cs;
-        cs.factory_name = "dsn::tools::asio_udp_provider";
-        cs.message_buffer_block_size = 1024 * 64;
-        spec.network_default_client_cfs[RPC_CHANNEL_UDP] = cs;
-    }
-    {
-        network_server_config cs2;
-        cs2.port = 0;
-        cs2.channel = RPC_CHANNEL_UDP;
-        cs2.factory_name = "dsn::tools::asio_udp_provider";
-        cs2.message_buffer_block_size = 1024 * 64;
-        spec.network_default_server_cfs[cs2] = cs2;
+    if (FLAGS_enable_udp) {
+        {
+            network_client_config client_conf;
+            client_conf.factory_name = "dsn::tools::asio_udp_provider";
+            client_conf.message_buffer_block_size = 1024 * 64;
+            spec.network_default_client_cfs[RPC_CHANNEL_UDP] = client_conf;
+        }
+        {
+            network_server_config server_conf;
+            server_conf.port = 0;
+            server_conf.channel = RPC_CHANNEL_UDP;
+            server_conf.factory_name = "dsn::tools::asio_udp_provider";
+            server_conf.message_buffer_block_size = 1024 * 64;
+            spec.network_default_server_cfs[server_conf] = server_conf;
+        }
     }
 
     if (spec.logging_factory_name == "")

diff --git a/src/runtime/providers.common.cpp b/src/runtime/providers.common.cpp
@@ -39,6 +39,7 @@
 #include "runtime/task/simple_task_queue.h"
 #include "runtime/task/hpc_task_queue.h"
 #include "runtime/rpc/network.sim.h"
+#include "utils/flags.h"
 #include "utils/simple_logger.h"
 #include "runtime/rpc/dsn_message_parser.h"
 #include "runtime/rpc/thrift_message_parser.h"
@@ -47,6 +48,8 @@
 namespace dsn {
 namespace tools {
 
+DSN_DEFINE_bool("network", enable_udp, true, "whether to enable udp rpc engine");
+
 void register_std_lock_providers()
 {
     lock_provider::register_component<std_lock_provider>("dsn::tools::std_lock_provider");
@@ -64,8 +67,10 @@ void register_common_providers()
 
     register_std_lock_providers();
 
+    if (FLAGS_enable_udp) {
+        register_component_provider<asio_udp_provider>("dsn::tools::asio_udp_provider");
+    }
     register_component_provider<asio_network_provider>("dsn::tools::asio_network_provider");
-    register_component_provider<asio_udp_provider>("dsn::tools::asio_udp_provider");
     register_component_provider<sim_network_provider>("dsn::tools::sim_network_provider");
     register_component_provider<simple_task_queue>("dsn::tools::simple_task_queue");
     register_component_provider<hpc_concurrent_task_queue>("dsn::tools::hpc_concurrent_task_queue");

diff --git a/src/runtime/task/task_spec.cpp b/src/runtime/task/task_spec.cpp
@@ -29,12 +29,16 @@
 #include <array>
 
 #include "runtime/rpc/rpc_message.h"
+#include "utils/flags.h"
 #include "utils/fmt_logging.h"
 #include "utils/command_manager.h"
 #include "utils/threadpool_spec.h"
 #include "utils/smart_pointers.h"
 
 namespace dsn {
+namespace tools {
+DSN_DECLARE_bool(enable_udp);
+}
 
 constexpr int TASK_SPEC_STORE_CAPACITY = 512;
 
@@ -224,6 +228,12 @@ bool task_spec::init()
                 return false;
             }
         }
+
+        if (spec->rpc_call_channel == RPC_CHANNEL_UDP && !dsn::tools::FLAGS_enable_udp) {
+            LOG_ERROR("task rpc_call_channel RPC_CHANNEL_UCP need udp service, make sure "
+                      "[network].enable_udp");
+            return false;
+        }
     }
 
     return true;