udp: support forwarding packets between workers

include/envoy/event/dispatcher.h

@@ -161,7 +161,7 @@ class Dispatcher {
    * @param cb supplies the udp listener callbacks to invoke for listener events.
    * @return Network::ListenerPtr a new listener that is owned by the caller.
    */
-  virtual Network::UdpListenerPtr createUdpListener(Network::SocketSharedPtr&& socket,
+  virtual Network::UdpListenerPtr createUdpListener(Network::SocketSharedPtr socket,
                                                     Network::UdpListenerCallbacks& cb) PURE;
   /**
    * Allocates a timer. @see Timer for docs on how to use the timer.

include/envoy/network/connection_handler.h

@@ -52,6 +52,13 @@ class ConnectionHandler {
    */
   virtual void removeListeners(uint64_t listener_tag) PURE;
 
+  /**
+   * Get the ``UdpListenerCallbacks`` associated with ``listener_tag``. This will be
+   * nullptr for non-UDP listeners and for ``listener_tag`` values that have already been
+   * removed.
+   */
+  virtual UdpListenerCallbacks* getUdpListenerCallbacks(uint64_t listener_tag) PURE;
+
   /**
    * Remove the filter chains and the connections in the listener. All connections owned
    * by the filter chains will be closed. Once all the connections are destroyed(connections
@@ -147,8 +154,8 @@ class ActiveUdpListenerFactory {
    * @return the ActiveUdpListener created.
    */
   virtual ConnectionHandler::ActiveListenerPtr
-  createActiveUdpListener(ConnectionHandler& parent, Event::Dispatcher& disptacher,
-                          Network::ListenerConfig& config) PURE;
+  createActiveUdpListener(uint32_t worker_id, ConnectionHandler& parent,
+                          Event::Dispatcher& dispatcher, Network::ListenerConfig& config) PURE;
 
   /**
    * @return true if the UDP passing through listener doesn't form stateful connections.
@@ -159,4 +166,4 @@ class ActiveUdpListenerFactory {
 using ActiveUdpListenerFactoryPtr = std::unique_ptr<ActiveUdpListenerFactory>;
 
 } // namespace Network
-} // namespace Envoy
+} // namespace Envoy

include/envoy/network/listener.h

@@ -20,6 +20,7 @@ namespace Envoy {
 namespace Network {
 
 class ActiveUdpListenerFactory;
+class UdpListenerWorkerRouter;
 
 /**
  * ListenSocketFactory is a member of ListenConfig to provide listen socket.
@@ -142,6 +143,12 @@ class ListenerConfig {
    */
   virtual UdpPacketWriterFactoryOptRef udpPacketWriterFactory() PURE;
 
+  /**
+   * @return the ``UdpListenerWorkerRouter`` for this listener. This will
+   * be non-null iff this is a UDP listener.
+   */
+  virtual UdpListenerWorkerRouter* udpListenerWorkerRouter() PURE;
+
   /**
    * @return traffic direction of the listener.
    */
@@ -246,7 +253,7 @@ class UdpListenerCallbacks {
    *
    * @param data UdpRecvData from the underlying socket.
    */
-  virtual void onData(UdpRecvData& data) PURE;
+  virtual void onData(UdpRecvData&& data) PURE;
 
   /**
    * Called when the underlying socket is ready for read, before onData() is
@@ -278,6 +285,31 @@ class UdpListenerCallbacks {
    * UdpListenerCallback
    */
   virtual UdpPacketWriter& udpPacketWriter() PURE;
+
+  /**
+   * Returns the id of this worker, in the range of [0, concurrency).
+   */
+  virtual uint32_t workerId() const PURE;
+
+  /**
+   * Called whenever data is received on the underlying udp socket, on
+   * the destination worker for the datagram according to ``destination()``.
+   */
+  virtual void onDataWorker(Network::UdpRecvData& data) PURE;
+
+  /**
+   * Returns the worker id that ``data`` should be delivered to. A return value
+   * of ``absl::nullopt`` indicates the packet should be processed on the current
+   * worker. The return value must be in the range [0, concurrency), or ``absl::nullopt``.
+   * @param concurrency specifies the concurrency, which is the upper bound for worker id.
+   */
+  virtual absl::optional<uint32_t> destination(const Network::UdpRecvData& data,
+                                               uint32_t concurrency) PURE;
+
+  /**
+   * Posts ``data`` to be delivered on this worker.
+   */
+  virtual void post(Network::UdpRecvData&& data) PURE;
 };
 
 /**
@@ -341,5 +373,32 @@ class UdpListener : public virtual Listener {
 
 using UdpListenerPtr = std::unique_ptr<UdpListener>;
 
+/**
+ * Handles delivering datagrams to the correct worker.
+ */
+class UdpListenerWorkerRouter {
+public:
+  virtual ~UdpListenerWorkerRouter() = default;
+
+  /**
+   * Registers a worker's callbacks for this listener. This worker must accept
+   * packets until it calls ``unregisterWorker``.
+   */
+  virtual void registerWorker(UdpListenerCallbacks& listener) PURE;
+
+  /**
+   * Unregisters a worker's callbacks for this listener.
+   */
+  virtual void unregisterWorker(UdpListenerCallbacks& listener) PURE;
+
+  /**
+   * Deliver ``data`` to the correct worker by calling ``onDataWorker()``
+   * or ``post()`` on one of the registered workers.
+   */
+  virtual void deliver(UdpListenerCallbacks& current, UdpRecvData&& data) PURE;
+};
+
+using UdpListenerWorkerRouterPtr = std::unique_ptr<UdpListenerWorkerRouter>;
+
 } // namespace Network
 } // namespace Envoy

include/envoy/server/worker.h

@@ -100,11 +100,12 @@ class WorkerFactory {
   virtual ~WorkerFactory() = default;
 
   /**
+   * @param id supplies the id of the worker, in the range of [0, concurrency).
    * @param overload_manager supplies the server's overload manager.
    * @param worker_name supplies the name of the worker, used for per-worker stats.
    * @return WorkerPtr a new worker.
    */
-  virtual WorkerPtr createWorker(OverloadManager& overload_manager,
+  virtual WorkerPtr createWorker(uint32_t id, OverloadManager& overload_manager,
                                  const std::string& worker_name) PURE;
 };
 

source/common/event/dispatcher_impl.cc

@@ -144,7 +144,7 @@ Network::ListenerPtr DispatcherImpl::createListener(Network::SocketSharedPtr&& s
                                                     backlog_size);
 }
 
-Network::UdpListenerPtr DispatcherImpl::createUdpListener(Network::SocketSharedPtr&& socket,
+Network::UdpListenerPtr DispatcherImpl::createUdpListener(Network::SocketSharedPtr socket,
                                                           Network::UdpListenerCallbacks& cb) {
   ASSERT(isThreadSafe());
   return std::make_unique<Network::UdpListenerImpl>(*this, std::move(socket), cb, timeSource());

source/common/event/dispatcher_impl.h

@@ -62,7 +62,7 @@ class DispatcherImpl : Logger::Loggable<Logger::Id::main>,
   Network::ListenerPtr createListener(Network::SocketSharedPtr&& socket,
                                       Network::TcpListenerCallbacks& cb, bool bind_to_port,
                                       uint32_t backlog_size) override;
-  Network::UdpListenerPtr createUdpListener(Network::SocketSharedPtr&& socket,
+  Network::UdpListenerPtr createUdpListener(Network::SocketSharedPtr socket,
                                             Network::UdpListenerCallbacks& cb) override;
   TimerPtr createTimer(TimerCb cb) override;
   Event::SchedulableCallbackPtr createSchedulableCallback(std::function<void()> cb) override;

source/common/network/udp_listener_impl.cc

@@ -93,7 +93,7 @@ void UdpListenerImpl::processPacket(Address::InstanceConstSharedPtr local_addres
   ASSERT(local_address != nullptr);
   UdpRecvData recvData{
       {std::move(local_address), std::move(peer_address)}, std::move(buffer), receive_time};
-  cb_.onData(recvData);
+  cb_.onData(std::move(recvData));
 }
 
 void UdpListenerImpl::handleWriteCallback() {
@@ -125,5 +125,47 @@ Api::IoCallUint64Result UdpListenerImpl::flush() {
   return cb_.udpPacketWriter().flush();
 }
 
+UdpListenerWorkerRouterImpl::UdpListenerWorkerRouterImpl(uint32_t concurrency)
+    : workers_(concurrency) {}
+
+void UdpListenerWorkerRouterImpl::registerWorker(UdpListenerCallbacks& listener) {
+  absl::WriterMutexLock lock(&mutex_);
+
+  ASSERT(listener.workerId() < workers_.size());
+  workers_.at(listener.workerId()) = &listener;
+}
+
+void UdpListenerWorkerRouterImpl::unregisterWorker(UdpListenerCallbacks& listener) {
+  absl::WriterMutexLock lock(&mutex_);
+
+  ASSERT(workers_.at(listener.workerId()) == &listener);
+  workers_.at(listener.workerId()) = nullptr;
+}
+
+void UdpListenerWorkerRouterImpl::deliver(UdpListenerCallbacks& current, UdpRecvData&& data) {
+  absl::ReaderMutexLock lock(&mutex_);
+
+  absl::optional<uint32_t> dest;
+
+  // For concurrency == 1, the packet will always go to the current worker.
+  if (workers_.size() > 1) {
+    dest = current.destination(data, workers_.size());
+  }
+
+  if (!dest.has_value() || *dest == current.workerId()) {
+    current.onDataWorker(data);
+  } else {
+    RELEASE_ASSERT(*dest < workers_.size(),
+                   "UdpListenerCallbacks::destination returned out-of-range value");
+    auto* worker = workers_[*dest];
+
+    // When a listener is being removed, packets could be processed on some workers after the
+    // listener is removed from other workers, which could result in a nullptr for that worker.
+    if (worker != nullptr) {
+      worker->post(std::move(data));
+    }
+  }
+}
+
 } // namespace Network
 } // namespace Envoy

source/common/network/udp_listener_impl.h

@@ -61,5 +61,18 @@ class UdpListenerImpl : public BaseListenerImpl,
   Event::FileEventPtr file_event_;
 };
 
+class UdpListenerWorkerRouterImpl : public UdpListenerWorkerRouter {
+public:
+  UdpListenerWorkerRouterImpl(uint32_t concurrency);
+
+  void registerWorker(UdpListenerCallbacks& listener) override;
+  void unregisterWorker(UdpListenerCallbacks& listener) override;
+  void deliver(UdpListenerCallbacks& current, UdpRecvData&& data) override;
+
+private:
+  absl::Mutex mutex_;
+  std::vector<UdpListenerCallbacks*> workers_ ABSL_GUARDED_BY(mutex_);
+};
+
 } // namespace Network
 } // namespace Envoy

source/common/runtime/runtime_features.cc

@@ -78,6 +78,7 @@ constexpr const char* runtime_features[] = {
     "envoy.reloadable_features.listener_in_place_filterchain_update",
     "envoy.reloadable_features.new_tcp_connection_pool",
     "envoy.reloadable_features.overload_manager_disable_keepalive_drain_http2",
+    "envoy.reloadable_features.prefer_quic_kernel_bpf_packet_routing",
     "envoy.reloadable_features.preserve_query_string_in_path_redirects",
     "envoy.reloadable_features.preserve_upstream_date",
     "envoy.reloadable_features.stop_faking_paths",

source/extensions/quic_listeners/quiche/BUILD

@@ -409,10 +409,7 @@ envoy_cc_library(
         "//bazel:linux": ["udp_gso_batch_writer.cc"],
         "//conditions:default": [],
     }),
-    hdrs = select({
-        "//bazel:linux": ["udp_gso_batch_writer.h"],
-        "//conditions:default": [],
-    }),
+    hdrs = ["udp_gso_batch_writer.h"],
     external_deps = ["quiche_quic_platform"],
     tags = ["nofips"],
     visibility = [