Internal change

PiperOrigin-RevId: 370194092

cpp/core/internal/mediums/webrtc.cc

@@ -62,15 +62,6 @@ WebRtc::~WebRtc() {
   for (const auto& service_id : service_ids) {
     StopAcceptingConnections(service_id);
   }
-
-  // Disconnect all connections
-  absl::flat_hash_set<std::string> peer_ids;
-  for (auto& item : sockets_) {
-    peer_ids.emplace(item.first);
-  }
-  for (const auto& peer_id : peer_ids) {
-    sockets_.find(peer_id)->second.Close();
-  }
 }
 
 const std::string WebRtc::GetDefaultCountryCode() {
@@ -653,28 +644,16 @@ void WebRtc::RestartTachyonReceiveMessages(const std::string& service_id) {
              service_id.c_str());
 }
 
-void WebRtc::ProcessDataChannelCreated(
-    const std::string& service_id, const PeerId& remote_peer_id,
-    rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
+void WebRtc::ProcessDataChannelOpen(const std::string& service_id,
+                                    const PeerId& remote_peer_id,
+                                    WebRtcSocketWrapper socket_wrapper) {
   MutexLock lock(&mutex_);
 
-  // Transform the DataChannel into a socket.
-  auto socket = std::make_unique<WebRtcSocket>("WebRtcSocket", data_channel);
-  socket->SetOnSocketClosedListener({[this, remote_peer_id]() {
-    OffloadFromThread("rtc-socket-closed-cb", [this, remote_peer_id]() {
-      ProcessDataChannelClosed(remote_peer_id);
-    });
-  }});
-  WebRtcSocketWrapper wrapper = WebRtcSocketWrapper(std::move(socket));
-
-  // Store this DataChannel so that we can update it later.
-  sockets_.emplace(remote_peer_id.GetId(), wrapper);
-
   // Notify the client of the newly formed socket.
   const auto& connection_request_entry =
       requesting_connections_info_.find(remote_peer_id.GetId());
   if (connection_request_entry != requesting_connections_info_.end()) {
-    connection_request_entry->second.socket_future.Set(wrapper);
+    connection_request_entry->second.socket_future.Set(socket_wrapper);
     return;
   }
 
@@ -682,52 +661,23 @@ void WebRtc::ProcessDataChannelCreated(
       accepting_connections_info_.find(service_id);
   if (accepting_connection_entry != accepting_connections_info_.end()) {
     accepting_connection_entry->second.accepted_connection_callback.accepted_cb(
-        wrapper);
+        socket_wrapper);
     return;
   }
 
   // No one to handle the newly created DataChannel, so we'll just close it.
-  wrapper.Close();
+  socket_wrapper.Close();
   NEARBY_LOG(INFO,
              "Ignoring new DataChannel because we "
              "are not accepting connections for service %s.",
              service_id.c_str());
 }
 
-void WebRtc::ProcessDataChannelMessage(const PeerId& remote_peer_id,
-                                       const ByteArray& message) {
-  MutexLock lock(&mutex_);
-  const auto& entry = sockets_.find(remote_peer_id.GetId());
-  if (entry == sockets_.end()) {
-    return;
-  }
-
-  entry->second.NotifyDataChannelMsgReceived(message);
-}
-
-void WebRtc::ProcessDataChannelBufferAmountChanged(
-    const PeerId& remote_peer_id) {
-  MutexLock lock(&mutex_);
-  const auto& entry = sockets_.find(remote_peer_id.GetId());
-  if (entry == sockets_.end()) {
-    return;
-  }
-
-  entry->second.NotifyDataChannelBufferedAmountChanged();
-}
-
 void WebRtc::ProcessDataChannelClosed(const PeerId& remote_peer_id) {
   MutexLock lock(&mutex_);
   NEARBY_LOG(INFO,
              "Data channel has closed, removing connection flow for peer %s.",
              remote_peer_id.GetId().c_str());
-  const auto& entry = sockets_.find(remote_peer_id.GetId());
-  if (entry == sockets_.end()) {
-    return;
-  }
-
-  entry->second.Close();
-  sockets_.erase(remote_peer_id.GetId());
 
   RemoveConnectionFlow(remote_peer_id);
 }
@@ -753,28 +703,13 @@ std::unique_ptr<ConnectionFlow> WebRtc::CreateConnectionFlow(
                  });
            }}},
       {
-          .data_channel_created_cb =
-              {[this, service_id,
-                remote_peer_id](rtc::scoped_refptr<webrtc::DataChannelInterface>
-                                    data_channel) {
-                OffloadFromThread(
-                    "rtc-channel-created",
-                    [this, service_id, remote_peer_id, data_channel]() {
-                      ProcessDataChannelCreated(service_id, remote_peer_id,
-                                                data_channel);
-                    });
-              }},
-          .data_channel_message_received_cb = {[this, remote_peer_id](
-                                                   const ByteArray& message) {
+          .data_channel_open_cb = {[this, service_id, remote_peer_id](
+                                       WebRtcSocketWrapper socket_wrapper) {
             OffloadFromThread(
-                "rtc-channel-messsage", [this, remote_peer_id, message]() {
-                  ProcessDataChannelMessage(remote_peer_id, message);
-                });
-          }},
-          .data_channel_buffered_amount_changed_cb = {[this, remote_peer_id]() {
-            OffloadFromThread(
-                "rtc-channel-buffer-change", [this, remote_peer_id]() {
-                  ProcessDataChannelBufferAmountChanged(remote_peer_id);
+                "rtc-channel-created",
+                [this, service_id, remote_peer_id, socket_wrapper]() {
+                  ProcessDataChannelOpen(service_id, remote_peer_id,
+                                         socket_wrapper);
                 });
           }},
           .data_channel_closed_cb = {[this, remote_peer_id]() {

cpp/core/internal/mediums/webrtc.h

@@ -209,18 +209,9 @@ class WebRtc {
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Runs on |single_thread_executor_|.
-  void ProcessDataChannelCreated(
-      const std::string& service_id, const PeerId& remote_peer_id,
-      rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel)
-      ABSL_LOCKS_EXCLUDED(mutex_);
-
-  // Runs on |single_thread_executor_|.
-  void ProcessDataChannelMessage(const PeerId& remote_peer_id,
-                                 const ByteArray& message)
-      ABSL_LOCKS_EXCLUDED(mutex_);
-
-  // Runs on |single_thread_executor_|.
-  void ProcessDataChannelBufferAmountChanged(const PeerId& remote_peer_id)
+  void ProcessDataChannelOpen(const std::string& service_id,
+                              const PeerId& remote_peer_id,
+                              WebRtcSocketWrapper socket_wrapper)
       ABSL_LOCKS_EXCLUDED(mutex_);
 
   // Runs on |single_thread_executor_|.
@@ -265,11 +256,6 @@ class WebRtc {
   // a unique ConnectionFlow.
   absl::flat_hash_map<std::string, std::unique_ptr<ConnectionFlow>>
       connection_flows_ ABSL_GUARDED_BY(mutex_);
-
-  // A map of a remote PeerId -> Socket. Non-empty while we have active
-  // connections.
-  absl::flat_hash_map<std::string, WebRtcSocketWrapper> sockets_
-      ABSL_GUARDED_BY(mutex_);
 };
 
 }  // namespace mediums

cpp/core/internal/mediums/webrtc/BUILD

@@ -16,7 +16,6 @@ cc_library(
     name = "webrtc",
     srcs = [
         "connection_flow.cc",
-        "data_channel_observer_impl.cc",
         "peer_id.cc",
         "signaling_frames.cc",
         "webrtc_socket.cc",
@@ -24,7 +23,6 @@ cc_library(
     hdrs = [
         "connection_flow.h",
         "data_channel_listener.h",
-        "data_channel_observer_impl.h",
         "local_ice_candidate_listener.h",
         "peer_id.h",
         "session_description_wrapper.h",

cpp/core/internal/mediums/webrtc/connection_flow.cc

@@ -18,7 +18,8 @@
 #include <memory>
 
 #include "core/internal/mediums/webrtc/session_description_wrapper.h"
-#include "platform/public/count_down_latch.h"
+#include "core/internal/mediums/webrtc/webrtc_socket.h"
+#include "core/internal/mediums/webrtc/webrtc_socket_wrapper.h"
 #include "platform/public/logging.h"
 #include "platform/public/mutex_lock.h"
 #include "platform/public/webrtc.h"
@@ -135,13 +136,9 @@ ConnectionFlow::ConnectionFlow(
 
 ConnectionFlow::~ConnectionFlow() {
   NEARBY_LOG(INFO, "~ConnectionFlow");
-  CountDownLatch latch(1);
-  if (RunOnSignalingThread([this, latch]() mutable {
-        CloseOnSignalingThread();
-        latch.CountDown();
-      })) {
-    latch.Await();
-  }
+  RunOnSignalingThread([this] { CloseOnSignalingThread(); });
+  shutdown_latch_.Await();
+  NEARBY_LOG(INFO, "~ConnectionFlow done");
 }
 
 SessionDescriptionWrapper ConnectionFlow::CreateOffer() {
@@ -170,9 +167,8 @@ void ConnectionFlow::CreateOfferOnSignalingThread(
   webrtc::DataChannelInit data_channel_init;
   data_channel_init.reliable = true;
   auto pc = GetPeerConnection();
-  rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel =
-      pc->CreateDataChannel(kDataChannelName, &data_channel_init);
-  RegisterDataChannelObserver(data_channel);
+  CreateSocketFromDataChannel(
+      pc->CreateDataChannel(kDataChannelName, &data_channel_init));
 
   webrtc::PeerConnectionInterface::RTCOfferAnswerOptions options;
   rtc::scoped_refptr<CreateSessionDescriptionObserverImpl> observer =
@@ -384,6 +380,7 @@ bool ConnectionFlow::InitPeerConnection(WebRtcMedium& webrtc_medium) {
   ExceptionOr<bool> result = success_future.Get(kPeerConnectionTimeout);
   bool success = result.ok() && result.result();
   if (!success) {
+    shutdown_latch_.CountDown();
     NEARBY_LOG(ERROR, "Failed to create peer connection: %d",
                result.exception());
   }
@@ -401,6 +398,31 @@ void ConnectionFlow::OnSignalingStable() {
   cached_remote_ice_candidates_.clear();
 }
 
+void ConnectionFlow::CreateSocketFromDataChannel(
+    rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
+  NEARBY_LOG(INFO, "Creating data channel socket");
+  auto socket =
+      std::make_unique<WebRtcSocket>("WebRtcSocket", std::move(data_channel));
+  socket->SetSocketListener({
+      .socket_ready_cb = {[this](WebRtcSocket* socket) {
+        CHECK(IsRunningOnSignalingThread());
+        if (!TransitionState(State::kWaitingToConnect, State::kConnected)) {
+          NEARBY_LOG(ERROR,
+                     "Data channel socket is open but connection flow was not "
+                     "in the required state");
+          socket->Close();
+          return;
+        }
+        // Pass socket wrapper by copy on purpose
+        data_channel_listener_.data_channel_open_cb(socket_wrapper_);
+      }},
+      .socket_closed_cb = [callback =
+                               data_channel_listener_.data_channel_closed_cb](
+                              WebRtcSocket*) { callback(); },
+  });
+  socket_wrapper_ = WebRtcSocketWrapper(std::move(socket));
+}
+
 void ConnectionFlow::OnIceCandidate(
     const webrtc::IceCandidateInterface* candidate) {
   CHECK(IsRunningOnSignalingThread());
@@ -420,7 +442,7 @@ void ConnectionFlow::OnDataChannel(
     rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
   NEARBY_LOG(INFO, "OnDataChannel");
   CHECK(IsRunningOnSignalingThread());
-  RegisterDataChannelObserver(std::move(data_channel));
+  CreateSocketFromDataChannel(std::move(data_channel));
 }
 
 void ConnectionFlow::OnIceGatheringChange(
@@ -447,37 +469,6 @@ void ConnectionFlow::OnRenegotiationNeeded() {
   CHECK(IsRunningOnSignalingThread());
 }
 
-void ConnectionFlow::ProcessDataChannelConnectedOnSignalingThread(
-    rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-  NEARBY_LOG(INFO, "Data channel state changed to connected.");
-  CHECK(IsRunningOnSignalingThread());
-  if (!TransitionState(State::kWaitingToConnect, State::kConnected)) {
-    data_channel->Close();
-    return;
-  }
-
-  data_channel_listener_.data_channel_created_cb(std::move(data_channel));
-}
-
-void ConnectionFlow::RegisterDataChannelObserver(
-    rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-  CHECK(IsRunningOnSignalingThread());
-  if (!data_channel_observer_) {
-    auto state_change_callback = [this, data_channel]() {
-      if (data_channel->state() ==
-          webrtc::DataChannelInterface::DataState::kOpen) {
-        ProcessDataChannelConnectedOnSignalingThread(std::move(data_channel));
-      }
-    };
-    data_channel_observer_ = absl::make_unique<DataChannelObserverImpl>(
-        data_channel, &data_channel_listener_,
-        std::move(state_change_callback));
-    NEARBY_LOG(INFO, "Registered data channel observer");
-  } else {
-    NEARBY_LOG(WARNING, "Data channel observer already exists");
-  }
-}
-
 bool ConnectionFlow::TransitionState(State current_state, State new_state) {
   CHECK(IsRunningOnSignalingThread());
   if (current_state != state_) {
@@ -497,13 +488,19 @@ bool ConnectionFlow::CloseOnSignalingThread() {
     return false;
   }
   state_ = State::kEnded;
+  // This prevents other tasks from queuing on the signaling thread for this
+  // object.
   auto pc = GetAndResetPeerConnection();
 
-  NEARBY_LOG(INFO, "Closing WebRTC connection.");
+  NEARBY_LOG(INFO, "Closing WebRTC peer connection.");
+  // NOTE: Closing the peer conection will close the data channel and thus the
+  // socket implicitly.
   if (pc) pc->Close();
-  data_channel_observer_.reset();
+  NEARBY_LOG(INFO, "Closed WebRTC peer connection.");
+  // Prevent any already queued tasks from running on the signaling thread
   can_run_tasks_.reset();
-  NEARBY_LOG(INFO, "Closed WebRTC connection.");
+  // If anyone was waiting for shutdown to be done let them know.
+  shutdown_latch_.CountDown();
   return true;
 }
 
@@ -525,7 +522,7 @@ bool ConnectionFlow::RunOnSignalingThread(Runnable&& runnable) {
         // (signaling thread). This guarantees that if the weak_ptr is valid
         // when this task starts, it will stay valid until the task ends.
         if (!can_run_tasks.lock()) {
-          NEARBY_LOG(VERBOSE,
+          NEARBY_LOG(INFO,
                      "Peer connection already closed. Cannot run tasks.");
           return;
         }

cpp/core/internal/mediums/webrtc/connection_flow.h

@@ -18,10 +18,11 @@
 #include <memory>
 
 #include "core/internal/mediums/webrtc/data_channel_listener.h"
-#include "core/internal/mediums/webrtc/data_channel_observer_impl.h"
 #include "core/internal/mediums/webrtc/local_ice_candidate_listener.h"
 #include "core/internal/mediums/webrtc/session_description_wrapper.h"
+#include "core/internal/mediums/webrtc/webrtc_socket_wrapper.h"
 #include "platform/base/runnable.h"
+#include "platform/public/count_down_latch.h"
 #include "platform/public/single_thread_executor.h"
 #include "platform/public/webrtc.h"
 #include "webrtc/api/data_channel_interface.h"
@@ -152,7 +153,8 @@ class ConnectionFlow : public webrtc::PeerConnectionObserver {
           ice_candidates);
   // Invoked when the peer connection indicates that signaling is stable.
   void OnSignalingStable() ABSL_LOCKS_EXCLUDED(mutex_);
-  void RegisterDataChannelObserver(
+
+  void CreateSocketFromDataChannel(
       rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel);
 
   // TODO(bfranz): Consider whether this needs to be configurable per platform
@@ -168,23 +170,21 @@ class ConnectionFlow : public webrtc::PeerConnectionObserver {
                                    State expected_entry_state,
                                    State exit_state);
 
-  void ProcessDataChannelConnectedOnSignalingThread(
-      rtc::scoped_refptr<webrtc::DataChannelInterface>)
-      ABSL_LOCKS_EXCLUDED(mutex_);
-
   bool CloseOnSignalingThread() ABSL_LOCKS_EXCLUDED(mutex_);
 
   bool RunOnSignalingThread(Runnable&& runnable);
   bool IsRunningOnSignalingThread();
 
   Mutex mutex_;
+  // Used to prevent the destructor from returning while the signaling thread is
+  // still running CloseOnSignalingThread()
+  CountDownLatch shutdown_latch_{1};
 
   // State is used on signaling thread only.
   State state_ = State::kInitialized;
+  // Used to communicate data channel events back to the caller of Create()
   DataChannelListener data_channel_listener_;
 
-  std::unique_ptr<DataChannelObserverImpl> data_channel_observer_;
-
   LocalIceCandidateListener local_ice_candidate_listener_;
   // Peer connection can be used only on signaling thread. The only exception
   // is accessing the signaling thread handle. Tasks posted on the
@@ -205,6 +205,10 @@ class ConnectionFlow : public webrtc::PeerConnectionObserver {
   rtc::scoped_refptr<webrtc::PeerConnectionInterface> peer_connection_
       ABSL_GUARDED_BY(mutex_);
 
+  // Used to hold a reference to the WebRtcSocket while the data channel is
+  // connecting.
+  WebRtcSocketWrapper socket_wrapper_;
+
   std::vector<std::unique_ptr<webrtc::IceCandidateInterface>>
       cached_remote_ice_candidates_;
   // This pointer is only for DCHECK() assertions.

cpp/core/internal/mediums/webrtc/connection_flow_test.cc

@@ -18,6 +18,7 @@
 #include <vector>
 
 #include "core/internal/mediums/webrtc/session_description_wrapper.h"
+#include "core/internal/mediums/webrtc/webrtc_socket_wrapper.h"
 #include "platform/base/byte_array.h"
 #include "platform/base/medium_environment.h"
 #include "platform/public/count_down_latch.h"
@@ -58,10 +59,7 @@ TEST_F(ConnectionFlowTest, SuccessfulOfferAnswerFlow) {
 
   Future<ByteArray> message_received_future;
 
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_data_channel_future;
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_data_channel_future;
+  Future<WebRtcSocketWrapper> offerer_socket_future, answerer_socket_future;
 
   std::unique_ptr<ConnectionFlow> offerer, answerer;
 
@@ -76,10 +74,9 @@ TEST_F(ConnectionFlowTest, SuccessfulOfferAnswerFlow) {
              if (answerer)
                answerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&offerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             offerer_data_channel_future.Set(std::move(data_channel));
+      {.data_channel_open_cb =
+           [&offerer_socket_future](WebRtcSocketWrapper socket) {
+             offerer_socket_future.Set(std::move(socket));
            }},
       webrtc_medium_offerer);
   ASSERT_NE(offerer, nullptr);
@@ -93,14 +90,9 @@ TEST_F(ConnectionFlowTest, SuccessfulOfferAnswerFlow) {
              if (offerer)
                offerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&answerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             answerer_data_channel_future.Set(std::move(data_channel));
-           },
-       .data_channel_message_received_cb =
-           [&message_received_future](ByteArray bytes) {
-             message_received_future.Set(std::move(bytes));
+      {.data_channel_open_cb =
+           [&answerer_socket_future](WebRtcSocketWrapper socket) {
+             answerer_socket_future.Set(std::move(socket));
            }},
       webrtc_medium_answerer);
   ASSERT_NE(answerer, nullptr);
@@ -118,18 +110,19 @@ TEST_F(ConnectionFlowTest, SuccessfulOfferAnswerFlow) {
   EXPECT_TRUE(answerer->SetLocalSessionDescription(std::move(answer)));
 
   // Retrieve Data Channels
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_channel = offerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(offerer_channel.ok());
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_channel = answerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(answerer_channel.ok());
+  ExceptionOr<WebRtcSocketWrapper> offerer_socket =
+      offerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(offerer_socket.ok());
+  ExceptionOr<WebRtcSocketWrapper> answerer_socket =
+      answerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(answerer_socket.ok());
 
   // Send message on data channel
   const char message[] = "Test";
-  offerer_channel.result()->Send(webrtc::DataBuffer(message));
+  offerer_socket.result().GetImpl().GetOutputStream().Write(
+      ByteArray(message, 4));
   ExceptionOr<ByteArray> received_message =
-      message_received_future.Get(absl::Seconds(1));
+      answerer_socket.result().GetImpl().GetInputStream().Read(4);
   EXPECT_TRUE(received_message.ok());
   EXPECT_EQ(received_message.result(), ByteArray{message});
 }
@@ -245,10 +238,7 @@ TEST_F(ConnectionFlowTest, TerminateAnswerer) {
 
   Future<ByteArray> message_received_future;
 
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_data_channel_future;
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_data_channel_future;
+  Future<WebRtcSocketWrapper> offerer_socket_future, answerer_socket_future;
 
   std::unique_ptr<ConnectionFlow> offerer, answerer;
 
@@ -263,10 +253,9 @@ TEST_F(ConnectionFlowTest, TerminateAnswerer) {
              if (answerer)
                answerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&offerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             offerer_data_channel_future.Set(std::move(data_channel));
+      {.data_channel_open_cb =
+           [&offerer_socket_future](WebRtcSocketWrapper socket) {
+             offerer_socket_future.Set(std::move(socket));
            }},
       webrtc_medium_offerer);
   ASSERT_NE(offerer, nullptr);
@@ -280,14 +269,9 @@ TEST_F(ConnectionFlowTest, TerminateAnswerer) {
              if (offerer)
                offerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&answerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             answerer_data_channel_future.Set(std::move(data_channel));
-           },
-       .data_channel_message_received_cb =
-           [&message_received_future](ByteArray bytes) {
-             message_received_future.Set(std::move(bytes));
+      {.data_channel_open_cb =
+           [&answerer_socket_future](WebRtcSocketWrapper wrapper) {
+             answerer_socket_future.Set(std::move(wrapper));
            }},
       webrtc_medium_answerer);
   ASSERT_NE(answerer, nullptr);
@@ -305,12 +289,12 @@ TEST_F(ConnectionFlowTest, TerminateAnswerer) {
   EXPECT_TRUE(answerer->SetLocalSessionDescription(std::move(answer)));
 
   // Retrieve Data Channels
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_channel = offerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(offerer_channel.ok());
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_channel = answerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(answerer_channel.ok());
+  ExceptionOr<WebRtcSocketWrapper> offerer_socket =
+      offerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(offerer_socket.ok());
+  ExceptionOr<WebRtcSocketWrapper> answerer_socket =
+      answerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(offerer_socket.ok());
 
   CountDownLatch latch(1);
   auto pc = answerer->GetPeerConnection();
@@ -321,10 +305,10 @@ TEST_F(ConnectionFlowTest, TerminateAnswerer) {
   latch.Await();
 
   // Send message on data channel
-  const char message[] = "Test";
-  offerer_channel.result()->Send(webrtc::DataBuffer(message));
+  std::string message = "Test";
+  offerer_socket.result().GetOutputStream().Write(ByteArray{message});
   ExceptionOr<ByteArray> received_message =
-      message_received_future.Get(absl::Seconds(1));
+      answerer_socket.result().GetInputStream().Read(4);
   EXPECT_FALSE(received_message.ok());
 }
 
@@ -333,10 +317,7 @@ TEST_F(ConnectionFlowTest, TerminateOfferer) {
 
   Future<ByteArray> message_received_future;
 
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_data_channel_future;
-  Future<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_data_channel_future;
+  Future<WebRtcSocketWrapper> offerer_socket_future, answerer_socket_future;
 
   std::unique_ptr<ConnectionFlow> offerer, answerer;
 
@@ -351,10 +332,9 @@ TEST_F(ConnectionFlowTest, TerminateOfferer) {
              if (answerer)
                answerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&offerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             offerer_data_channel_future.Set(std::move(data_channel));
+      {.data_channel_open_cb =
+           [&offerer_socket_future](WebRtcSocketWrapper socket) {
+             offerer_socket_future.Set(std::move(socket));
            }},
       webrtc_medium_offerer);
   ASSERT_NE(offerer, nullptr);
@@ -368,14 +348,9 @@ TEST_F(ConnectionFlowTest, TerminateOfferer) {
              if (offerer)
                offerer->OnRemoteIceCandidatesReceived(std::move(vec));
            }},
-      {.data_channel_created_cb =
-           [&answerer_data_channel_future](
-               rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel) {
-             answerer_data_channel_future.Set(std::move(data_channel));
-           },
-       .data_channel_message_received_cb =
-           [&message_received_future](ByteArray bytes) {
-             message_received_future.Set(std::move(bytes));
+      {.data_channel_open_cb =
+           [&answerer_socket_future](WebRtcSocketWrapper wrapper) {
+             answerer_socket_future.Set(std::move(wrapper));
            }},
       webrtc_medium_answerer);
   ASSERT_NE(answerer, nullptr);
@@ -393,12 +368,12 @@ TEST_F(ConnectionFlowTest, TerminateOfferer) {
   EXPECT_TRUE(answerer->SetLocalSessionDescription(std::move(answer)));
 
   // Retrieve Data Channels
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      offerer_channel = offerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(offerer_channel.ok());
-  ExceptionOr<rtc::scoped_refptr<webrtc::DataChannelInterface>>
-      answerer_channel = answerer_data_channel_future.Get(absl::Seconds(1));
-  EXPECT_TRUE(answerer_channel.ok());
+  ExceptionOr<WebRtcSocketWrapper> offerer_socket =
+      offerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(offerer_socket.ok());
+  ExceptionOr<WebRtcSocketWrapper> answerer_socket =
+      answerer_socket_future.Get(absl::Seconds(1));
+  EXPECT_TRUE(offerer_socket.ok());
 
   CountDownLatch latch(1);
   auto pc = offerer->GetPeerConnection();
@@ -409,10 +384,10 @@ TEST_F(ConnectionFlowTest, TerminateOfferer) {
   latch.Await();
 
   // Send message on data channel
-  const char message[] = "Test";
-  offerer_channel.result()->Send(webrtc::DataBuffer(message));
+  std::string message = "Test";
+  offerer_socket.result().GetOutputStream().Write(ByteArray{message});
   ExceptionOr<ByteArray> received_message =
-      message_received_future.Get(absl::Seconds(1));
+      answerer_socket.result().GetInputStream().Read(4);
   EXPECT_FALSE(received_message.ok());
 }
 }  // namespace

cpp/core/internal/mediums/webrtc/data_channel_listener.h

@@ -15,6 +15,7 @@
 #ifndef CORE_INTERNAL_MEDIUMS_WEBRTC_DATA_CHANNEL_LISTENER_H_
 #define CORE_INTERNAL_MEDIUMS_WEBRTC_DATA_CHANNEL_LISTENER_H_
 
+#include "core/internal/mediums/webrtc/webrtc_socket_wrapper.h"
 #include "core/listeners.h"
 #include "platform/base/byte_array.h"
 
@@ -25,19 +26,10 @@ namespace mediums {
 
 // Callbacks from the data channel.
 struct DataChannelListener {
-  // Called when the data channel is created.
-  std::function<void(rtc::scoped_refptr<webrtc::DataChannelInterface>)>
-      data_channel_created_cb =
-          DefaultCallback<rtc::scoped_refptr<webrtc::DataChannelInterface>>();
-
-  // Called when a new message was received on the data channel.
-  std::function<void(const ByteArray&)> data_channel_message_received_cb =
-      DefaultCallback<const ByteArray&>();
-
-  // Called when the data channel indicates that the buffered amount has
-  // changed.
-  std::function<void()> data_channel_buffered_amount_changed_cb =
-      DefaultCallback<>();
+  // Called when the data channel is open and the socket wraper is ready to
+  // read and write.
+  std::function<void(WebRtcSocketWrapper)> data_channel_open_cb =
+      DefaultCallback<WebRtcSocketWrapper>();
 
   // Called when the data channel is closed.
   std::function<void()> data_channel_closed_cb = DefaultCallback<>();

cpp/core/internal/mediums/webrtc/data_channel_observer_impl.cc

@@ -1,62 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// Licensed 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
-//
-//     https://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.
-
-#include "core/internal/mediums/webrtc/data_channel_observer_impl.h"
-
-namespace location {
-namespace nearby {
-namespace connections {
-namespace mediums {
-
-DataChannelObserverImpl::DataChannelObserverImpl(
-    rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel,
-    DataChannelListener* data_channel_listener,
-    DataChannelStateChangeCallback callback)
-    : data_channel_listener_(data_channel_listener),
-      state_change_callback_(std::move(callback)),
-      data_channel_{std::move(data_channel)} {
-  data_channel_->RegisterObserver(this);
-}
-
-DataChannelObserverImpl::~DataChannelObserverImpl() { Disconnect(); }
-
-void DataChannelObserverImpl::OnStateChange() {
-  if (data_channel_->state() ==
-      webrtc::DataChannelInterface::DataState::kClosed) {
-    Disconnect();
-  }
-  state_change_callback_();
-}
-
-void DataChannelObserverImpl::OnMessage(const webrtc::DataBuffer& buffer) {
-  data_channel_listener_->data_channel_message_received_cb(
-      ByteArray(buffer.data.data<char>(), buffer.size()));
-}
-
-void DataChannelObserverImpl::OnBufferedAmountChange(uint64_t sent_data_size) {
-  data_channel_listener_->data_channel_buffered_amount_changed_cb();
-}
-
-void DataChannelObserverImpl::Disconnect() {
-  data_channel_->UnregisterObserver();
-  if (data_channel_listener_) {
-    data_channel_listener_->data_channel_closed_cb();
-    data_channel_listener_ = nullptr;
-  }
-}
-
-}  // namespace mediums
-}  // namespace connections
-}  // namespace nearby
-}  // namespace location

cpp/core/internal/mediums/webrtc/data_channel_observer_impl.h

@@ -1,57 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// Licensed 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
-//
-//     https://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.
-
-#ifndef CORE_INTERNAL_MEDIUMS_WEBRTC_DATA_CHANNEL_OBSERVER_IMPL_H_
-#define CORE_INTERNAL_MEDIUMS_WEBRTC_DATA_CHANNEL_OBSERVER_IMPL_H_
-
-#include "core/internal/mediums/webrtc/data_channel_listener.h"
-#include "webrtc/api/data_channel_interface.h"
-
-namespace location {
-namespace nearby {
-namespace connections {
-namespace mediums {
-
-class DataChannelObserverImpl : public webrtc::DataChannelObserver {
- public:
-  using DataChannelStateChangeCallback = std::function<void()>;
-
-  // Creates and registers an observer for |data_channel|
-  // The observer is unregistered in destructor or when |data_channel|
-  // is closed.
-  DataChannelObserverImpl(
-      rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel,
-      DataChannelListener* data_channel_listener,
-      DataChannelStateChangeCallback callback);
-  ~DataChannelObserverImpl() override;
-
-  // webrtc::DataChannelObserver:
-  void OnStateChange() override;
-  void OnMessage(const webrtc::DataBuffer& buffer) override;
-  void OnBufferedAmountChange(uint64_t sent_data_size) override;
-
- private:
-  void Disconnect();
-
-  DataChannelListener* data_channel_listener_;
-  DataChannelStateChangeCallback state_change_callback_;
-  rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel_;
-};
-
-}  // namespace mediums
-}  // namespace connections
-}  // namespace nearby
-}  // namespace location
-
-#endif  // CORE_INTERNAL_MEDIUMS_WEBRTC_DATA_CHANNEL_OBSERVER_IMPL_H_

cpp/core/internal/mediums/webrtc/webrtc_socket.cc

@@ -57,32 +57,89 @@ Exception WebRtcSocket::OutputStreamImpl::Close() {
 WebRtcSocket::WebRtcSocket(
     const std::string& name,
     rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel)
-    : name_(name), data_channel_(std::move(data_channel)) {}
+    : name_(name), data_channel_(std::move(data_channel)) {
+  NEARBY_LOGS(INFO) << "WebRtcSocket::WebRtcSocket(" << name_
+                       << ") this: " << this;
+  data_channel_->RegisterObserver(this);
+}
+
+WebRtcSocket::~WebRtcSocket() {
+  NEARBY_LOGS(INFO) << "WebRtcSocket::~WebRtcSocket(" << name_
+                       << ") this: " << this;
+  Close();
+  NEARBY_LOGS(INFO) << "WebRtcSocket::~WebRtcSocket(" << name_
+                       << ") this: " << this << " done";
+}
 
 InputStream& WebRtcSocket::GetInputStream() { return pipe_.GetInputStream(); }
 
 OutputStream& WebRtcSocket::GetOutputStream() { return output_stream_; }
 
 void WebRtcSocket::Close() {
+  NEARBY_LOGS(INFO) << "WebRtcSocket::Close(" << name_ << ") this: " << this;
   if (closed_.Set(true)) return;
 
-  pipe_.GetInputStream().Close();
-  pipe_.GetOutputStream().Close();
+  ClosePipe();
+  // NOTE: This call blocks and triggers a state change on the siginaling thread
+  // to 'closing' but does not block until 'closed' is sent so the data channel
+  // is not fully closed when this call is done.
   data_channel_->Close();
-  WakeUpWriter();
-  socket_closed_listener_.socket_closed_cb();
+  NEARBY_LOGS(INFO) << "WebRtcSocket::Close(" << name_ << ") this: " << this
+                       << " done";
 }
 
-void WebRtcSocket::NotifyDataChannelMsgReceived(const ByteArray& message) {
-  if (!pipe_.GetOutputStream().Write(message).Ok()) {
-    Close();
-    return;
+void WebRtcSocket::OnStateChange() {
+  // Running on the signaling thread right now.
+  NEARBY_LOGS(ERROR)
+      << "WebRtcSocket::OnStateChange() webrtc data channel state: "
+      << webrtc::DataChannelInterface::DataStateString(data_channel_->state());
+  switch (data_channel_->state()) {
+    case webrtc::DataChannelInterface::DataState::kConnecting:
+      break;
+    case webrtc::DataChannelInterface::DataState::kOpen:
+      // We implicitly depend on the |socket_listener_| to offload from
+      // the signaling thread so it does not get blocked.
+      socket_listener_.socket_ready_cb(this);
+      break;
+    case webrtc::DataChannelInterface::DataState::kClosing:
+      break;
+    case webrtc::DataChannelInterface::DataState::kClosed:
+      NEARBY_LOG(
+          ERROR,
+          "WebRtcSocket::OnStateChange() unregistering data channel observer.");
+      data_channel_->UnregisterObserver();
+      // This will trigger a destruction of the owning connection flow
+      // We implicitly depend on the |socket_listener_| to offload from
+      // the signaling thread so it does not get blocked.
+      socket_listener_.socket_closed_cb(this);
+
+      if (!closed_.Set(true)) {
+        ClosePipe();
+      }
+      break;
   }
+}
+void WebRtcSocket::OnMessage(const webrtc::DataBuffer& buffer) {
+  // This is a data channel callback on the signaling thread, lets off load so
+  // we don't block signaling.
+  OffloadFromSignalingThread(
+      [this, buffer = ByteArray(buffer.data.data<char>(), buffer.size())] {
+        if (!pipe_.GetOutputStream().Write(buffer).Ok()) {
+          Close();
+          return;
+        }
 
-  if (!pipe_.GetOutputStream().Flush().Ok()) Close();
+        if (!pipe_.GetOutputStream().Flush().Ok()) {
+          Close();
+        }
+      });
 }
 
-void WebRtcSocket::NotifyDataChannelBufferedAmountChanged() { WakeUpWriter(); }
+void WebRtcSocket::OnBufferedAmountChange(uint64_t sent_data_size) {
+  // This is a data channel callback on the signaling thread, lets off load so
+  // we don't block signaling.
+  OffloadFromSignalingThread([this] { WakeUpWriter(); });
+}
 
 bool WebRtcSocket::SendMessage(const ByteArray& data) {
   return data_channel_->Send(
@@ -91,13 +148,26 @@ bool WebRtcSocket::SendMessage(const ByteArray& data) {
 
 bool WebRtcSocket::IsClosed() { return closed_.Get(); }
 
+void WebRtcSocket::ClosePipe() {
+  NEARBY_LOGS(INFO) << "WebRtcSocket::ClosePipe(" << name_
+                       << ") this: " << this;
+  // This is thread-safe to close these sockets even if a read or write is in
+  // process on another thread, Close will wait for the exclusive mutex before
+  // setting state.
+  pipe_.GetInputStream().Close();
+  pipe_.GetOutputStream().Close();
+  WakeUpWriter();
+  NEARBY_LOGS(INFO) << "WebRtcSocket::ClosePipe(" << name_
+                       << ") this: " << this << " done";
+}
+
 void WebRtcSocket::WakeUpWriter() {
   MutexLock lock(&backpressure_mutex_);
   buffer_variable_.Notify();
 }
 
-void WebRtcSocket::SetOnSocketClosedListener(SocketClosedListener&& listener) {
-  socket_closed_listener_ = std::move(listener);
+void WebRtcSocket::SetSocketListener(SocketListener&& listener) {
+  socket_listener_ = std::move(listener);
 }
 
 void WebRtcSocket::BlockUntilSufficientSpaceInBuffer(int length) {
@@ -109,6 +179,10 @@ void WebRtcSocket::BlockUntilSufficientSpaceInBuffer(int length) {
   }
 }
 
+void WebRtcSocket::OffloadFromSignalingThread(Runnable runnable) {
+  single_thread_executor_.Execute(std::move(runnable));
+}
+
 }  // namespace mediums
 }  // namespace connections
 }  // namespace nearby

cpp/core/internal/mediums/webrtc/webrtc_socket.h

@@ -25,7 +25,9 @@
 #include "platform/public/condition_variable.h"
 #include "platform/public/mutex.h"
 #include "platform/public/pipe.h"
+#include "platform/public/single_thread_executor.h"
 #include "webrtc/api/data_channel_interface.h"
+
 namespace location {
 namespace nearby {
 namespace connections {
@@ -39,11 +41,11 @@ constexpr int kMaxDataSize = 1 * 1024 * 1024;
 //
 // Messages are buffered here to prevent the data channel from overflowing,
 // which could lead to data loss.
-class WebRtcSocket : public Socket {
+class WebRtcSocket : public Socket, public webrtc::DataChannelObserver {
  public:
   WebRtcSocket(const std::string& name,
                rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel);
-  ~WebRtcSocket() override = default;
+  ~WebRtcSocket() override;
 
   WebRtcSocket(const WebRtcSocket& other) = delete;
   WebRtcSocket& operator=(const WebRtcSocket& other) = delete;
@@ -53,20 +55,20 @@ class WebRtcSocket : public Socket {
   OutputStream& GetOutputStream() override;
   void Close() override;
 
-  // Callback from WebRTC data channel when new message has been received from
-  // the remote.
-  void NotifyDataChannelMsgReceived(const ByteArray& message);
+  // webrtc::DataChannelObserver:
+  void OnStateChange() override;
+  void OnMessage(const webrtc::DataBuffer& buffer) override;
+  void OnBufferedAmountChange(uint64_t sent_data_size) override;
 
-  // Callback from WebRTC data channel that the buffered data amount has
-  // changed.
-  void NotifyDataChannelBufferedAmountChanged();
-
-  // Listener class the gets called when the socket is closed.
-  struct SocketClosedListener {
-    std::function<void()> socket_closed_cb = DefaultCallback<>();
+  // Listener class the gets called when the socket is ready or closed
+  struct SocketListener {
+    std::function<void(WebRtcSocket*)> socket_ready_cb =
+        DefaultCallback<WebRtcSocket*>();
+    std::function<void(WebRtcSocket*)> socket_closed_cb =
+        DefaultCallback<WebRtcSocket*>();
   };
 
-  void SetOnSocketClosedListener(SocketClosedListener&& listener);
+  void SetSocketListener(SocketListener&& listener);
 
  private:
   class OutputStreamImpl : public OutputStream {
@@ -89,8 +91,10 @@ class WebRtcSocket : public Socket {
 
   void WakeUpWriter();
   bool IsClosed();
+  void ClosePipe();
   bool SendMessage(const ByteArray& data);
   void BlockUntilSufficientSpaceInBuffer(int length);
+  void OffloadFromSignalingThread(Runnable runnable);
 
   std::string name_;
   rtc::scoped_refptr<webrtc::DataChannelInterface> data_channel_;
@@ -101,10 +105,14 @@ class WebRtcSocket : public Socket {
 
   AtomicBoolean closed_{false};
 
-  SocketClosedListener socket_closed_listener_;
+  SocketListener socket_listener_;
 
   mutable Mutex backpressure_mutex_;
   ConditionVariable buffer_variable_{&backpressure_mutex_};
+
+  // This should be destroyed first to ensure any remaining tasks flushed on
+  // shutdown get run while the other members are still alive.
+  SingleThreadExecutor single_thread_executor_;
 };
 
 }  // namespace mediums

cpp/core/internal/mediums/webrtc/webrtc_socket_test.cc

@@ -58,23 +58,23 @@ class MockDataChannel
 }  // namespace
 
 TEST(WebRtcSocketTest, ReadFromSocket) {
-  const ByteArray kMessage{"Message"};
+  const char* message = "message";
   rtc::scoped_refptr<MockDataChannel> mock_data_channel = new MockDataChannel();
   WebRtcSocket webrtc_socket(kSocketName, mock_data_channel);
 
-  webrtc_socket.NotifyDataChannelMsgReceived(kMessage);
+  webrtc_socket.OnMessage(webrtc::DataBuffer{message});
   ExceptionOr<ByteArray> result = webrtc_socket.GetInputStream().Read(7);
   EXPECT_TRUE(result.ok());
-  EXPECT_EQ(result.result(), kMessage);
+  EXPECT_EQ(result.result(), ByteArray{message});
 }
 
 TEST(WebRtcSocketTest, ReadMultipleMessages) {
   rtc::scoped_refptr<MockDataChannel> mock_data_channel = new MockDataChannel();
   WebRtcSocket webrtc_socket(kSocketName, mock_data_channel);
 
-  webrtc_socket.NotifyDataChannelMsgReceived(ByteArray{"Me"});
-  webrtc_socket.NotifyDataChannelMsgReceived(ByteArray{"ssa"});
-  webrtc_socket.NotifyDataChannelMsgReceived(ByteArray{"ge"});
+  webrtc_socket.OnMessage(webrtc::DataBuffer{"Me"});
+  webrtc_socket.OnMessage(webrtc::DataBuffer{"ssa"});
+  webrtc_socket.OnMessage(webrtc::DataBuffer{"ge"});
 
   ExceptionOr<ByteArray> result;
 
@@ -131,10 +131,19 @@ TEST(WebRtcSocketTest, Close) {
 
   int socket_closed_cb_called = 0;
 
-  webrtc_socket.SetOnSocketClosedListener(
-      {.socket_closed_cb = [&]() { socket_closed_cb_called++; }});
+  webrtc_socket.SetSocketListener(
+      {.socket_closed_cb = [&](WebRtcSocket* socket) {
+        socket_closed_cb_called++;
+      }});
   webrtc_socket.Close();
 
+  // We have to fake the close event to get the callback to run.
+  ON_CALL(*mock_data_channel, state())
+      .WillByDefault(
+          testing::Return(webrtc::DataChannelInterface::DataState::kClosed));
+
+  webrtc_socket.OnStateChange();
+
   EXPECT_EQ(socket_closed_cb_called, 1);
 }
 
@@ -162,6 +171,63 @@ TEST(WebRtcSocketTest, ReadFromClosedChannel) {
   EXPECT_EQ(webrtc_socket.GetInputStream().Read(7).exception(), Exception::kIo);
 }
 
+TEST(WebRtcSocketTest, DataChannelCloseEventCleansUp) {
+  rtc::scoped_refptr<MockDataChannel> mock_data_channel = new MockDataChannel();
+  WebRtcSocket webrtc_socket(kSocketName, mock_data_channel);
+
+  ON_CALL(*mock_data_channel, state())
+      .WillByDefault(
+          testing::Return(webrtc::DataChannelInterface::DataState::kClosed));
+
+  webrtc_socket.OnStateChange();
+
+  EXPECT_EQ(webrtc_socket.GetInputStream().Read(7).exception(), Exception::kIo);
+
+  // Calling Close again should be safe even if the channel is already shut
+  // down.
+  webrtc_socket.Close();
+}
+
+TEST(WebRtcSocketTest, OpenStateTriggersCallback) {
+  rtc::scoped_refptr<MockDataChannel> mock_data_channel = new MockDataChannel();
+  WebRtcSocket webrtc_socket(kSocketName, mock_data_channel);
+
+  int socket_ready_cb_called = 0;
+
+  webrtc_socket.SetSocketListener(
+      {.socket_ready_cb = [&](WebRtcSocket* socket) {
+        socket_ready_cb_called++;
+      }});
+
+  ON_CALL(*mock_data_channel, state())
+      .WillByDefault(
+          testing::Return(webrtc::DataChannelInterface::DataState::kOpen));
+
+  webrtc_socket.OnStateChange();
+
+  EXPECT_EQ(socket_ready_cb_called, 1);
+}
+
+TEST(WebRtcSocketTest, CloseStateTriggersCallback) {
+  rtc::scoped_refptr<MockDataChannel> mock_data_channel = new MockDataChannel();
+  WebRtcSocket webrtc_socket(kSocketName, mock_data_channel);
+
+  int socket_closed_cb_called = 0;
+
+  webrtc_socket.SetSocketListener(
+      {.socket_closed_cb = [&](WebRtcSocket* socket) {
+        socket_closed_cb_called++;
+      }});
+
+  ON_CALL(*mock_data_channel, state())
+      .WillByDefault(
+          testing::Return(webrtc::DataChannelInterface::DataState::kClosed));
+
+  webrtc_socket.OnStateChange();
+
+  EXPECT_EQ(socket_closed_cb_called, 1);
+}
+
 }  // namespace mediums
 }  // namespace connections
 }  // namespace nearby

cpp/core/internal/mediums/webrtc/webrtc_socket_wrapper.h

@@ -37,14 +37,6 @@ class WebRtcSocketWrapper final {
 
   OutputStream& GetOutputStream() { return impl_->GetOutputStream(); }
 
-  void NotifyDataChannelMsgReceived(const ByteArray& message) {
-    impl_->NotifyDataChannelMsgReceived(message);
-  }
-
-  void NotifyDataChannelBufferedAmountChanged() {
-    impl_->NotifyDataChannelBufferedAmountChanged();
-  }
-
   void Close() { return impl_->Close(); }
 
   bool IsValid() const { return impl_ != nullptr; }