Remove RepeatingFunction, use Alarm instead.

Remove RepeatingFunction from platform/api/task_runner.h, since:

1. It's not an API the embedder must implement.
2. There was a TODO to move it to util/.
3. Alarm was landed at the same time, and satisfies all the use cases,
   and doesn't need WeakPtrs to be safe; so, just have everything use
   that instead.

This change also fixes a lingering bug w.r.t. using invalid iterators
in QuicClient and QuicServer, which corrupted the std::map's tracking
the connections.

Bug: openscreen:77
Change-Id: If0ff4c90d85e083d11f5b18c140134541f8508a9
Reviewed-on: https://chromium-review.googlesource.com/c/openscreen/+/1932224
Reviewed-by: Ryan Keane <rwkeane@google.com>
Reviewed-by: Brandon Tolsch <btolsch@chromium.org>
Commit-Queue: Yuri Wiitala <miu@chromium.org>

17 files changed, 105 insertions, 77 deletions

diff --git a/osp/impl/discovery/mdns/mdns_responder_adapter.h b/osp/impl/discovery/mdns/mdns_responder_adapter.h
index 09092842..54fb440f 100644
--- a/osp/impl/discovery/mdns/mdns_responder_adapter.h
+++ b/osp/impl/discovery/mdns/mdns_responder_adapter.h
@@ -10,7 +10,6 @@
 #include <string>
 #include <vector>
 
-#include "absl/types/optional.h"
 #include "osp/impl/discovery/mdns/domain_name.h"
 #include "osp/impl/discovery/mdns/mdns_responder_platform.h"
 #include "platform/api/network_interface.h"
@@ -192,7 +191,7 @@ class MdnsResponderAdapter : public platform::UdpSocket::Client {
 
   // Returns the time period after which this method must be called again, if
   // any.
-  virtual absl::optional<platform::Clock::duration> RunTasks() = 0;
+  virtual platform::Clock::duration RunTasks() = 0;
 
   virtual std::vector<PtrEvent> TakePtrResponses() = 0;
   virtual std::vector<SrvEvent> TakeSrvResponses() = 0;
diff --git a/osp/impl/discovery/mdns/mdns_responder_adapter_impl.cc b/osp/impl/discovery/mdns/mdns_responder_adapter_impl.cc
index e79f3f1f..971b8cbf 100644
--- a/osp/impl/discovery/mdns/mdns_responder_adapter_impl.cc
+++ b/osp/impl/discovery/mdns/mdns_responder_adapter_impl.cc
@@ -353,7 +353,7 @@ void MdnsResponderAdapterImpl::OnError(platform::UdpSocket* socket,
   OSP_UNIMPLEMENTED();
 }
 
-absl::optional<platform::Clock::duration> MdnsResponderAdapterImpl::RunTasks() {
+platform::Clock::duration MdnsResponderAdapterImpl::RunTasks() {
   TRACE_SCOPED(TraceCategory::mDNS, "MdnsResponderAdapterImpl::RunTasks");
 
   mDNS_Execute(&mdns_);
diff --git a/osp/impl/discovery/mdns/mdns_responder_adapter_impl.h b/osp/impl/discovery/mdns/mdns_responder_adapter_impl.h
index fc3c9bcd..f1f0028c 100644
--- a/osp/impl/discovery/mdns/mdns_responder_adapter_impl.h
+++ b/osp/impl/discovery/mdns/mdns_responder_adapter_impl.h
@@ -39,7 +39,7 @@ class MdnsResponderAdapterImpl final : public MdnsResponderAdapter {
   void OnSendError(platform::UdpSocket* socket, Error error) override;
   void OnError(platform::UdpSocket* socket, Error error) override;
 
-  absl::optional<platform::Clock::duration> RunTasks() override;
+  platform::Clock::duration RunTasks() override;
 
   std::vector<PtrEvent> TakePtrResponses() override;
   std::vector<SrvEvent> TakeSrvResponses() override;
diff --git a/osp/impl/internal_services.cc b/osp/impl/internal_services.cc
index b1428945..ce0d83d4 100644
--- a/osp/impl/internal_services.cc
+++ b/osp/impl/internal_services.cc
@@ -170,8 +170,10 @@ void InternalServices::InternalPlatformLinkage::DeregisterInterfaces(
   }
 }
 
-InternalServices::InternalServices(platform::TaskRunner* task_runner)
-    : mdns_service_(task_runner,
+InternalServices::InternalServices(platform::ClockNowFunctionPtr now_function,
+                                   platform::TaskRunner* task_runner)
+    : mdns_service_(now_function,
+                    task_runner,
                     kServiceName,
                     kServiceProtocol,
                     std::make_unique<MdnsResponderAdapterImplFactory>(),
@@ -196,7 +198,7 @@ InternalServices* InternalServices::ReferenceSingleton(
     platform::TaskRunner* task_runner) {
   if (!g_instance) {
     OSP_CHECK_EQ(g_instance_ref_count, 0);
-    g_instance = new InternalServices(task_runner);
+    g_instance = new InternalServices(&platform::Clock::now, task_runner);
   }
   ++g_instance_ref_count;
   return g_instance;
diff --git a/osp/impl/internal_services.h b/osp/impl/internal_services.h
index 7f3105b4..a8f7ee8c 100644
--- a/osp/impl/internal_services.h
+++ b/osp/impl/internal_services.h
@@ -18,6 +18,7 @@
 #include "osp/public/protocol_connection_client.h"
 #include "osp/public/protocol_connection_server.h"
 #include "platform/api/network_interface.h"
+#include "platform/api/time.h"
 #include "platform/api/udp_socket.h"
 #include "platform/base/ip_address.h"
 #include "platform/base/macros.h"
@@ -67,7 +68,8 @@ class InternalServices : platform::UdpSocket::Client {
 
   // The TaskRunner provided here should live for the duration of this
   // InternalService object's lifetime.
-  explicit InternalServices(platform::TaskRunner* task_runner);
+  InternalServices(platform::ClockNowFunctionPtr now_function,
+                   platform::TaskRunner* task_runner);
   ~InternalServices() override;
 
   void RegisterMdnsSocket(platform::UdpSocket* socket);
diff --git a/osp/impl/mdns_responder_service.cc b/osp/impl/mdns_responder_service.cc
index 5dc42bc0..b980362c 100644
--- a/osp/impl/mdns_responder_service.cc
+++ b/osp/impl/mdns_responder_service.cc
@@ -32,6 +32,7 @@ std::string ServiceIdFromServiceInstanceName(
 }  // namespace
 
 MdnsResponderService::MdnsResponderService(
+    platform::ClockNowFunctionPtr now_function,
     platform::TaskRunner* task_runner,
     const std::string& service_name,
     const std::string& service_protocol,
@@ -40,7 +41,8 @@ MdnsResponderService::MdnsResponderService(
     : service_type_{{service_name, service_protocol}},
       mdns_responder_factory_(std::move(mdns_responder_factory)),
       platform_(std::move(platform)),
-      task_runner_(task_runner) {}
+      task_runner_(task_runner),
+      background_tasks_alarm_(now_function, task_runner) {}
 
 MdnsResponderService::~MdnsResponderService() = default;
 
@@ -132,12 +134,7 @@ void MdnsResponderService::StartListenerInternal() {
 
   StartListening();
   ServiceListenerImpl::Delegate::SetState(ServiceListener::State::kRunning);
-  // TODO(rwkeane): Use new Alarm class instead once owning CL is merged in.
-  // Then it can be more effectively cancelled when the state changes away from
-  // 'running'.
-  platform::RepeatingFunction::Post(
-      task_runner_,
-      std::bind(&mdns::MdnsResponderAdapter::RunTasks, mdns_responder_.get()));
+  RunBackgroundTasks();
 }
 
 void MdnsResponderService::StartAndSuspendListenerInternal() {
@@ -171,14 +168,13 @@ void MdnsResponderService::SearchNowInternal(ServiceListener::State from) {
 }
 
 void MdnsResponderService::StartPublisherInternal() {
-  if (!mdns_responder_)
+  if (!mdns_responder_) {
     mdns_responder_ = mdns_responder_factory_->Create();
+  }
 
   StartService();
   ServicePublisherImpl::Delegate::SetState(ServicePublisher::State::kRunning);
-  platform::RepeatingFunction::Post(
-      task_runner_,
-      std::bind(&mdns::MdnsResponderAdapter::RunTasks, mdns_responder_.get()));
+  RunBackgroundTasks();
 }
 
 void MdnsResponderService::StartAndSuspendPublisherInternal() {
@@ -255,7 +251,7 @@ void MdnsResponderService::HandleMdnsEvents() {
     if (events_possible) {
       // NOTE: This still needs to be called here, even though it runs in the
       // background regularly, because we just finished processing MDNS events.
-      mdns_responder_->RunTasks();
+      RunBackgroundTasks();
     }
   } while (events_possible);
 
@@ -658,4 +654,13 @@ MdnsResponderService::GetNetworkInterfaceIndexFromSocket(
   return it->interface_info.index;
 }
 
+void MdnsResponderService::RunBackgroundTasks() {
+  if (!mdns_responder_) {
+    return;
+  }
+  const auto delay_until_next_run = mdns_responder_->RunTasks();
+  background_tasks_alarm_.ScheduleFromNow([this] { RunBackgroundTasks(); },
+                                          delay_until_next_run);
+}
+
 }  // namespace openscreen
diff --git a/osp/impl/mdns_responder_service.h b/osp/impl/mdns_responder_service.h
index 97a52da1..f5eb13f5 100644
--- a/osp/impl/mdns_responder_service.h
+++ b/osp/impl/mdns_responder_service.h
@@ -18,7 +18,9 @@
 #include "osp/impl/service_publisher_impl.h"
 #include "platform/api/network_interface.h"
 #include "platform/api/task_runner.h"
+#include "platform/api/time.h"
 #include "platform/base/ip_address.h"
+#include "util/alarm.h"
 
 namespace openscreen {
 
@@ -34,6 +36,7 @@ class MdnsResponderService : public ServiceListenerImpl::Delegate,
                              public platform::UdpSocket::Client {
  public:
   MdnsResponderService(
+      platform::ClockNowFunctionPtr now_function,
       platform::TaskRunner* task_runner,
       const std::string& service_name,
       const std::string& service_protocol,
@@ -160,6 +163,9 @@ class MdnsResponderService : public ServiceListenerImpl::Delegate,
   platform::NetworkInterfaceIndex GetNetworkInterfaceIndexFromSocket(
       const platform::UdpSocket* socket) const;
 
+  // Runs background tasks to manage the internal mDNS state.
+  void RunBackgroundTasks();
+
   // Service type separated as service name and service protocol for both
   // listening and publishing (e.g. {"_openscreen", "_udp"}).
   std::array<std::string, 2> service_type_;
@@ -195,7 +201,10 @@ class MdnsResponderService : public ServiceListenerImpl::Delegate,
 
   std::map<std::string, ServiceInfo> receiver_info_;
 
-  platform::TaskRunner* task_runner_;
+  platform::TaskRunner* const task_runner_;
+
+  // Scheduled to run periodic background tasks.
+  Alarm background_tasks_alarm_;
 
   friend class TestingMdnsResponderService;
 };
diff --git a/osp/impl/mdns_responder_service_unittest.cc b/osp/impl/mdns_responder_service_unittest.cc
index d01eb3d3..ff903466 100644
--- a/osp/impl/mdns_responder_service_unittest.cc
+++ b/osp/impl/mdns_responder_service_unittest.cc
@@ -28,7 +28,8 @@ class TestingMdnsResponderService final : public MdnsResponderService {
       const std::string& service_protocol,
       std::unique_ptr<MdnsResponderAdapterFactory> mdns_responder_factory,
       std::unique_ptr<MdnsPlatformService> platform_service)
-      : MdnsResponderService(task_runner,
+      : MdnsResponderService(&platform::FakeClock::now,
+                             task_runner,
                              service_name,
                              service_protocol,
                              std::move(mdns_responder_factory),
diff --git a/osp/impl/protocol_connection_client_factory.cc b/osp/impl/protocol_connection_client_factory.cc
index 46b26e38..43357210 100644
--- a/osp/impl/protocol_connection_client_factory.cc
+++ b/osp/impl/protocol_connection_client_factory.cc
@@ -10,6 +10,7 @@
 #include "osp/impl/quic/quic_connection_factory_impl.h"
 #include "osp/public/network_service_manager.h"
 #include "platform/api/task_runner.h"
+#include "platform/api/time.h"
 
 namespace openscreen {
 
@@ -21,7 +22,7 @@ ProtocolConnectionClientFactory::Create(
     platform::TaskRunner* task_runner) {
   return std::make_unique<QuicClient>(
       demuxer, std::make_unique<QuicConnectionFactoryImpl>(task_runner),
-      observer, task_runner);
+      observer, &platform::Clock::now, task_runner);
 }
 
 }  // namespace openscreen
diff --git a/osp/impl/protocol_connection_server_factory.cc b/osp/impl/protocol_connection_server_factory.cc
index 887c2ecd..78ab6f5b 100644
--- a/osp/impl/protocol_connection_server_factory.cc
+++ b/osp/impl/protocol_connection_server_factory.cc
@@ -10,6 +10,7 @@
 #include "osp/impl/quic/quic_server.h"
 #include "osp/public/network_service_manager.h"
 #include "platform/api/task_runner.h"
+#include "platform/api/time.h"
 
 namespace openscreen {
 
@@ -22,7 +23,7 @@ ProtocolConnectionServerFactory::Create(
     platform::TaskRunner* task_runner) {
   return std::make_unique<QuicServer>(
       config, demuxer, std::make_unique<QuicConnectionFactoryImpl>(task_runner),
-      observer, task_runner);
+      observer, &platform::Clock::now, task_runner);
 }
 
 }  // namespace openscreen
diff --git a/osp/impl/quic/quic_client.cc b/osp/impl/quic/quic_client.cc
index 24a066f3..3fb4f739 100644
--- a/osp/impl/quic/quic_client.cc
+++ b/osp/impl/quic/quic_client.cc
@@ -8,7 +8,6 @@
 #include <functional>
 #include <memory>
 
-#include "absl/types/optional.h"
 #include "platform/api/task_runner.h"
 #include "platform/api/time.h"
 #include "util/logging.h"
@@ -19,14 +18,11 @@ QuicClient::QuicClient(
     MessageDemuxer* demuxer,
     std::unique_ptr<QuicConnectionFactory> connection_factory,
     ProtocolConnectionServiceObserver* observer,
+    platform::ClockNowFunctionPtr now_function,
     platform::TaskRunner* task_runner)
     : ProtocolConnectionClient(demuxer, observer),
-      connection_factory_(std::move(connection_factory)) {
-  if (task_runner != nullptr) {
-    platform::RepeatingFunction::Post(task_runner,
-                                      std::bind(&QuicClient::Cleanup, this));
-  }
-}
+      connection_factory_(std::move(connection_factory)),
+      cleanup_alarm_(now_function, task_runner) {}
 
 QuicClient::~QuicClient() {
   CloseAllConnections();
@@ -36,6 +32,7 @@ bool QuicClient::Start() {
   if (state_ == State::kRunning)
     return false;
   state_ = State::kRunning;
+  Cleanup();  // Start periodic clean-ups.
   observer_->OnRunning();
   return true;
 }
@@ -45,27 +42,31 @@ bool QuicClient::Stop() {
     return false;
   CloseAllConnections();
   state_ = State::kStopped;
+  Cleanup();  // Final clean-up.
   observer_->OnStopped();
   return true;
 }
 
-absl::optional<platform::Clock::duration> QuicClient::Cleanup() {
+void QuicClient::Cleanup() {
   for (auto& entry : connections_) {
     entry.second.delegate->DestroyClosedStreams();
     if (!entry.second.delegate->has_streams())
       entry.second.connection->Close();
   }
 
-  for (auto& entry : delete_connections_)
-    connections_.erase(entry);
-
+  for (uint64_t endpoint_id : delete_connections_) {
+    auto it = connections_.find(endpoint_id);
+    if (it != connections_.end()) {
+      connections_.erase(it);
+    }
+  }
   delete_connections_.clear();
 
-  constexpr platform::Clock::duration kQuicCleanupFrequency =
+  constexpr platform::Clock::duration kQuicCleanupPeriod =
       std::chrono::milliseconds(500);
-  return state_ == State::kStopped
-             ? absl::optional<platform::Clock::duration>(absl::nullopt)
-             : absl::optional<platform::Clock::duration>(kQuicCleanupFrequency);
+  if (state_ != State::kStopped) {
+    cleanup_alarm_.ScheduleFromNow([this] { Cleanup(); }, kQuicCleanupPeriod);
+  }
 }
 
 QuicClient::ConnectRequest QuicClient::Connect(
@@ -148,8 +149,7 @@ void QuicClient::OnConnectionClosed(uint64_t endpoint_id,
   auto connection_entry = connections_.find(endpoint_id);
   if (connection_entry == connections_.end())
     return;
-
-  delete_connections_.emplace_back(connection_entry);
+  delete_connections_.push_back(endpoint_id);
 
   // TODO(crbug.com/openscreen/42): If we reset request IDs when a connection is
   // closed, we might end up re-using request IDs when a new connection is
diff --git a/osp/impl/quic/quic_client.h b/osp/impl/quic/quic_client.h
index b192dbe4..854e7eab 100644
--- a/osp/impl/quic/quic_client.h
+++ b/osp/impl/quic/quic_client.h
@@ -16,6 +16,7 @@
 #include "platform/api/task_runner.h"
 #include "platform/api/time.h"
 #include "platform/base/ip_address.h"
+#include "util/alarm.h"
 
 namespace openscreen {
 
@@ -41,6 +42,7 @@ class QuicClient final : public ProtocolConnectionClient,
   QuicClient(MessageDemuxer* demuxer,
              std::unique_ptr<QuicConnectionFactory> connection_factory,
              ProtocolConnectionServiceObserver* observer,
+             platform::ClockNowFunctionPtr now_function,
              platform::TaskRunner* task_runner);
   ~QuicClient() override;
 
@@ -94,7 +96,7 @@ class QuicClient final : public ProtocolConnectionClient,
 
   // Deletes dead QUIC connections then returns the time interval before this
   // method should be run again.
-  absl::optional<platform::Clock::duration> Cleanup();
+  void Cleanup();
 
   std::unique_ptr<QuicConnectionFactory> connection_factory_;
 
@@ -123,9 +125,12 @@ class QuicClient final : public ProtocolConnectionClient,
   // completed the QUIC handshake.
   std::map<uint64_t, ServiceConnectionData> connections_;
 
-  // Connections that need to be destroyed, but have to wait for the next event
-  // loop due to the underlying QUIC implementation's way of referencing them.
-  std::vector<decltype(connections_)::iterator> delete_connections_;
+  // Connections (endpoint IDs) that need to be destroyed, but have to wait for
+  // the next event loop due to the underlying QUIC implementation's way of
+  // referencing them.
+  std::vector<uint64_t> delete_connections_;
+
+  Alarm cleanup_alarm_;
 };
 
 }  // namespace openscreen
diff --git a/osp/impl/quic/quic_server.cc b/osp/impl/quic/quic_server.cc
index 1fe226a3..8e47ca10 100644
--- a/osp/impl/quic/quic_server.cc
+++ b/osp/impl/quic/quic_server.cc
@@ -7,7 +7,6 @@
 #include <functional>
 #include <memory>
 
-#include "absl/types/optional.h"
 #include "platform/api/task_runner.h"
 #include "platform/api/time.h"
 #include "util/logging.h"
@@ -19,15 +18,12 @@ QuicServer::QuicServer(
     MessageDemuxer* demuxer,
     std::unique_ptr<QuicConnectionFactory> connection_factory,
     ProtocolConnectionServer::Observer* observer,
+    platform::ClockNowFunctionPtr now_function,
     platform::TaskRunner* task_runner)
     : ProtocolConnectionServer(demuxer, observer),
       connection_endpoints_(config.connection_endpoints),
-      connection_factory_(std::move(connection_factory)) {
-  if (task_runner != nullptr) {
-    platform::RepeatingFunction::Post(task_runner,
-                                      std::bind(&QuicServer::Cleanup, this));
-  }
-}
+      connection_factory_(std::move(connection_factory)),
+      cleanup_alarm_(now_function, task_runner) {}
 
 QuicServer::~QuicServer() {
   CloseAllConnections();
@@ -38,6 +34,7 @@ bool QuicServer::Start() {
     return false;
   state_ = State::kRunning;
   connection_factory_->SetServerDelegate(this, connection_endpoints_);
+  Cleanup();  // Start periodic clean-ups.
   observer_->OnRunning();
   return true;
 }
@@ -48,6 +45,7 @@ bool QuicServer::Stop() {
   connection_factory_->SetServerDelegate(nullptr, {});
   CloseAllConnections();
   state_ = State::kStopped;
+  Cleanup();  // Final clean-up.
   observer_->OnStopped();
   return true;
 }
@@ -69,20 +67,23 @@ bool QuicServer::Resume() {
   return true;
 }
 
-absl::optional<platform::Clock::duration> QuicServer::Cleanup() {
+void QuicServer::Cleanup() {
   for (auto& entry : connections_)
     entry.second.delegate->DestroyClosedStreams();
 
-  for (auto& entry : delete_connections_)
-    connections_.erase(entry);
-
+  for (uint64_t endpoint_id : delete_connections_) {
+    auto it = connections_.find(endpoint_id);
+    if (it != connections_.end()) {
+      connections_.erase(it);
+    }
+  }
   delete_connections_.clear();
 
-  constexpr platform::Clock::duration kQuicCleanupFrequency =
+  constexpr platform::Clock::duration kQuicCleanupPeriod =
       std::chrono::milliseconds(500);
-  return state_ == State::kStopped
-             ? absl::optional<platform::Clock::duration>(absl::nullopt)
-             : absl::optional<platform::Clock::duration>(kQuicCleanupFrequency);
+  if (state_ != State::kStopped) {
+    cleanup_alarm_.ScheduleFromNow([this] { Cleanup(); }, kQuicCleanupPeriod);
+  }
 }
 
 std::unique_ptr<ProtocolConnection> QuicServer::CreateProtocolConnection(
@@ -138,8 +139,7 @@ void QuicServer::OnConnectionClosed(uint64_t endpoint_id,
   auto connection_entry = connections_.find(endpoint_id);
   if (connection_entry == connections_.end())
     return;
-
-  delete_connections_.emplace_back(connection_entry);
+  delete_connections_.push_back(endpoint_id);
 
   // TODO(crbug.com/openscreen/42): If we reset request IDs when a connection is
   // closed, we might end up re-using request IDs when a new connection is
diff --git a/osp/impl/quic/quic_server.h b/osp/impl/quic/quic_server.h
index 964a0ff2..77e03900 100644
--- a/osp/impl/quic/quic_server.h
+++ b/osp/impl/quic/quic_server.h
@@ -15,6 +15,7 @@
 #include "platform/api/task_runner.h"
 #include "platform/api/time.h"
 #include "platform/base/ip_address.h"
+#include "util/alarm.h"
 
 namespace openscreen {
 
@@ -35,6 +36,7 @@ class QuicServer final : public ProtocolConnectionServer,
              MessageDemuxer* demuxer,
              std::unique_ptr<QuicConnectionFactory> connection_factory,
              ProtocolConnectionServer::Observer* observer,
+             platform::ClockNowFunctionPtr now_function,
              platform::TaskRunner* task_runner);
   ~QuicServer() override;
 
@@ -72,7 +74,7 @@ class QuicServer final : public ProtocolConnectionServer,
 
   // Deletes dead QUIC connections then returns the time interval before this
   // method should be run again.
-  absl::optional<platform::Clock::duration> Cleanup();
+  void Cleanup();
 
   const std::vector<IPEndpoint> connection_endpoints_;
   std::unique_ptr<QuicConnectionFactory> connection_factory_;
@@ -95,9 +97,12 @@ class QuicServer final : public ProtocolConnectionServer,
   // completed the QUIC handshake.
   std::map<uint64_t, ServiceConnectionData> connections_;
 
-  // Connections that need to be destroyed, but have to wait for the next event
-  // loop due to the underlying QUIC implementation's way of referencing them.
-  std::vector<decltype(connections_)::iterator> delete_connections_;
+  // Connections (endpoint IDs) that need to be destroyed, but have to wait for
+  // the next event loop due to the underlying QUIC implementation's way of
+  // referencing them.
+  std::vector<uint64_t> delete_connections_;
+
+  Alarm cleanup_alarm_;
 };
 
 }  // namespace openscreen
diff --git a/osp/impl/quic/testing/quic_test_support.cc b/osp/impl/quic/testing/quic_test_support.cc
index c1223986..87bb867c 100644
--- a/osp/impl/quic/testing/quic_test_support.cc
+++ b/osp/impl/quic/testing/quic_test_support.cc
@@ -29,11 +29,9 @@ FakeQuicBridge::FakeQuicBridge(platform::FakeTaskRunner* task_runner,
   client_socket_ = std::make_unique<platform::FakeUdpSocket>(
       task_runner_, fake_client_factory.get());
 
-  // TODO(rwkeane): Pass actual task runner instead of nullptr once the fake
-  // task runner correctly respects the time delay for delayed tasks.
-  quic_client = std::make_unique<QuicClient>(controller_demuxer.get(),
-                                             std::move(fake_client_factory),
-                                             &mock_client_observer, nullptr);
+  quic_client = std::make_unique<QuicClient>(
+      controller_demuxer.get(), std::move(fake_client_factory),
+      &mock_client_observer, now_function, task_runner);
 
   auto fake_server_factory =
       std::make_unique<FakeServerQuicConnectionFactory>(fake_bridge.get());
@@ -41,9 +39,9 @@ FakeQuicBridge::FakeQuicBridge(platform::FakeTaskRunner* task_runner,
       task_runner_, fake_server_factory.get());
   ServerConfig config;
   config.connection_endpoints.push_back(kReceiverEndpoint);
-  quic_server = std::make_unique<QuicServer>(config, receiver_demuxer.get(),
-                                             std::move(fake_server_factory),
-                                             &mock_server_observer, nullptr);
+  quic_server = std::make_unique<QuicServer>(
+      config, receiver_demuxer.get(), std::move(fake_server_factory),
+      &mock_server_observer, now_function, task_runner);
 
   quic_client->Start();
   quic_server->Start();
diff --git a/osp/impl/testing/fake_mdns_responder_adapter.cc b/osp/impl/testing/fake_mdns_responder_adapter.cc
index a0989ec8..f4a9f4ae 100644
--- a/osp/impl/testing/fake_mdns_responder_adapter.cc
+++ b/osp/impl/testing/fake_mdns_responder_adapter.cc
@@ -251,8 +251,8 @@ void FakeMdnsResponderAdapter::OnError(platform::UdpSocket* socket,
   OSP_NOTREACHED() << "Tests should not drive this class with packets";
 }
 
-absl::optional<platform::Clock::duration> FakeMdnsResponderAdapter::RunTasks() {
-  return absl::nullopt;
+platform::Clock::duration FakeMdnsResponderAdapter::RunTasks() {
+  return std::chrono::seconds(1);
 }
 
 std::vector<mdns::PtrEvent> FakeMdnsResponderAdapter::TakePtrResponses() {
diff --git a/osp/impl/testing/fake_mdns_responder_adapter.h b/osp/impl/testing/fake_mdns_responder_adapter.h
index f6d54e6c..f1d2c1aa 100644
--- a/osp/impl/testing/fake_mdns_responder_adapter.h
+++ b/osp/impl/testing/fake_mdns_responder_adapter.h
@@ -116,7 +116,7 @@ class FakeMdnsResponderAdapter final : public mdns::MdnsResponderAdapter {
                           platform::UdpSocket* socket) override;
   Error DeregisterInterface(platform::UdpSocket* socket) override;
 
-  absl::optional<platform::Clock::duration> RunTasks() override;
+  platform::Clock::duration RunTasks() override;
 
   std::vector<mdns::PtrEvent> TakePtrResponses() override;
   std::vector<mdns::SrvEvent> TakeSrvResponses() override;