Use MaybeManagedPtr for QuickSocket for optional lifetime management

Summary:
`connSetupCallback_` and `connCallback_` in `QuicTransportBase.h` are raw pointers, which delegates the responsibility to keep these callbacks alive to the caller.

There are use cases where it would be convenient to be able to tie the lifetime of the callback to the Quic transport, e.g,. as long as the Quic transport is alive, it keeps the callbacks alive as well.

This diff uses MaybeManagedPtr to achieve this lifetime tie if desired. A MaybeManagedPtr intialized with a shared pointer manages lifetime of the contained object, whereas a MaybeManagedPtr initialized with a raw pointer does not manage lifetime of the contained object. This way caller can decide to pass in a shared ptr or raw pointer and achieve the desired behavior.

Note that we cannot simply use a shared_ptr for that. Using a shared_ptr would potentially mean that callbacks passed are destroyed when the transport is destroyed. Callbacks would not be destroyed if they were managed by a shared_ptr already, but this is something we cannot assume for every case. This would thus be a change in semantics to the current implementation, where the callbacks can outlive the transport.

Reviewed By: mjoras

Differential Revision: D49502381

fbshipit-source-id: 771a9328b99dc4f94f8e9679f9caf98af9180428

quic/api/QuicSocket.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <folly/Expected.h>
+#include <folly/MaybeManagedPtr.h>
 #include <folly/Optional.h>
 #include <folly/io/IOBuf.h>
 #include <folly/io/async/AsyncTransportCertificate.h>
@@ -273,13 +274,14 @@ class QuicSocket {
    * socket.
    */
   virtual void setConnectionSetupCallback(
-      ConnectionSetupCallback* callback) = 0;
+      folly::MaybeManagedPtr<ConnectionSetupCallback> callback) = 0;
 
   /**
    * Sets connection streams callback. This callback must be set after
    * connection set up is finished and is ready for streams processing.
    */
-  virtual void setConnectionCallback(ConnectionCallback* callback) = 0;
+  virtual void setConnectionCallback(
+      folly::MaybeManagedPtr<ConnectionCallback> callback) = 0;
 
   /**
    * Sets the functions that mvfst will invoke to validate early data params

quic/api/QuicTransportBase.cpp

@@ -1606,7 +1606,7 @@ void QuicTransportBase::logStreamOpenEvent(StreamId streamId) {
 void QuicTransportBase::handleNewStreams(std::vector<StreamId>& streamStorage) {
   const auto& newPeerStreamIds = streamStorage;
   for (const auto& streamId : newPeerStreamIds) {
-    CHECK_NOTNULL(connCallback_);
+    CHECK_NOTNULL(connCallback_.get());
     if (isBidirectionalStream(streamId)) {
       connCallback_->onNewBidirectionalStream(streamId);
     } else {
@@ -1625,7 +1625,7 @@ void QuicTransportBase::handleNewGroupedStreams(
     std::vector<StreamId>& streamStorage) {
   const auto& newPeerStreamIds = streamStorage;
   for (const auto& streamId : newPeerStreamIds) {
-    CHECK_NOTNULL(connCallback_);
+    CHECK_NOTNULL(connCallback_.get());
     auto stream = CHECK_NOTNULL(conn_->streamManager->getStream(streamId));
     CHECK(stream->groupId);
     if (isBidirectionalStream(streamId)) {
@@ -2845,11 +2845,12 @@ void QuicTransportBase::setAckRxTimestampsDisabled(
 }
 
 void QuicTransportBase::setConnectionSetupCallback(
-    ConnectionSetupCallback* callback) {
+    folly::MaybeManagedPtr<ConnectionSetupCallback> callback) {
   connSetupCallback_ = callback;
 }
 
-void QuicTransportBase::setConnectionCallback(ConnectionCallback* callback) {
+void QuicTransportBase::setConnectionCallback(
+    folly::MaybeManagedPtr<ConnectionCallback> callback) {
   connCallback_ = callback;
 }
 

quic/api/QuicTransportBase.h

@@ -241,9 +241,11 @@ class QuicTransportBase : public QuicSocket, QuicStreamPrioritiesObserver {
 
   virtual void setAckRxTimestampsDisabled(bool disableAckRxTimestamps);
 
-  void setConnectionSetupCallback(ConnectionSetupCallback* callback) final;
+  void setConnectionSetupCallback(
+      folly::MaybeManagedPtr<ConnectionSetupCallback> callback) final;
 
-  void setConnectionCallback(ConnectionCallback* callback) final;
+  void setConnectionCallback(
+      folly::MaybeManagedPtr<ConnectionCallback> callback) final;
 
   void setEarlyDataAppParamsFunctions(
       folly::Function<bool(const folly::Optional<std::string>&, const Buf&)
@@ -856,8 +858,8 @@ class QuicTransportBase : public QuicSocket, QuicStreamPrioritiesObserver {
 
   std::atomic<QuicEventBase*> qEvbPtr_;
   std::unique_ptr<QuicAsyncUDPSocketWrapper> socket_;
-  ConnectionSetupCallback* connSetupCallback_{nullptr};
-  ConnectionCallback* connCallback_{nullptr};
+  folly::MaybeManagedPtr<ConnectionSetupCallback> connSetupCallback_{nullptr};
+  folly::MaybeManagedPtr<ConnectionCallback> connCallback_{nullptr};
   // A flag telling transport if the new onConnectionEnd(error) cb must be used.
   bool useConnectionEndWithErrorCallback_{false};
 

quic/api/test/MockQuicSocket.h

@@ -135,8 +135,14 @@ class MockQuicSocket : public QuicSocket {
       (folly::Expected<folly::Unit, LocalErrorCode>),
       setReadCallback,
       (StreamId, ReadCallback*, folly::Optional<ApplicationErrorCode> err));
-  MOCK_METHOD(void, setConnectionSetupCallback, (ConnectionSetupCallback*));
-  MOCK_METHOD(void, setConnectionCallback, (ConnectionCallback*));
+  MOCK_METHOD(
+      void,
+      setConnectionSetupCallback,
+      (folly::MaybeManagedPtr<ConnectionSetupCallback>));
+  MOCK_METHOD(
+      void,
+      setConnectionCallback,
+      (folly::MaybeManagedPtr<ConnectionCallback>));
   void setEarlyDataAppParamsFunctions(
       folly::Function<bool(const folly::Optional<std::string>&, const Buf&)
                           const> validator,

quic/server/QuicServerTransport.cpp

@@ -25,8 +25,8 @@ namespace quic {
 QuicServerTransport::QuicServerTransport(
     folly::EventBase* evb,
     std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-    ConnectionSetupCallback* connSetupCb,
-    ConnectionCallback* connStreamsCb,
+    folly::MaybeManagedPtr<ConnectionSetupCallback> connSetupCb,
+    folly::MaybeManagedPtr<ConnectionCallback> connStreamsCb,
     std::shared_ptr<const fizz::server::FizzServerContext> ctx,
     std::unique_ptr<CryptoFactory> cryptoFactory,
     PacketNum startingPacketNum)
@@ -43,8 +43,8 @@ QuicServerTransport::QuicServerTransport(
 QuicServerTransport::QuicServerTransport(
     folly::EventBase* evb,
     std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-    ConnectionSetupCallback* connSetupCb,
-    ConnectionCallback* connStreamsCb,
+    folly::MaybeManagedPtr<ConnectionSetupCallback> connSetupCb,
+    folly::MaybeManagedPtr<ConnectionCallback> connStreamsCb,
     std::shared_ptr<const fizz::server::FizzServerContext> ctx,
     std::unique_ptr<CryptoFactory> cryptoFactory,
     bool useConnectionEndWithErrorCallback)
@@ -85,8 +85,8 @@ QuicServerTransport::~QuicServerTransport() {
 QuicServerTransport::Ptr QuicServerTransport::make(
     folly::EventBase* evb,
     std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-    ConnectionSetupCallback* connSetupCb,
-    ConnectionCallback* connStreamsCb,
+    const folly::MaybeManagedPtr<ConnectionSetupCallback>& connSetupCb,
+    const folly::MaybeManagedPtr<ConnectionCallback>& connStreamsCb,
     std::shared_ptr<const fizz::server::FizzServerContext> ctx,
     bool useConnectionEndWithErrorCallback) {
   return std::make_shared<QuicServerTransport>(

quic/server/QuicServerTransport.h

@@ -74,16 +74,16 @@ class QuicServerTransport
   static QuicServerTransport::Ptr make(
       folly::EventBase* evb,
       std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-      ConnectionSetupCallback* connSetupCb,
-      ConnectionCallback* connStreamsCb,
+      const folly::MaybeManagedPtr<ConnectionSetupCallback>& connSetupCb,
+      const folly::MaybeManagedPtr<ConnectionCallback>& connStreamsCb,
       std::shared_ptr<const fizz::server::FizzServerContext> ctx,
       bool useConnectionEndWithErrorCallback = false);
 
   QuicServerTransport(
       folly::EventBase* evb,
       std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-      ConnectionSetupCallback* connSetupCb,
-      ConnectionCallback* connStreamsCb,
+      folly::MaybeManagedPtr<ConnectionSetupCallback> connSetupCb,
+      folly::MaybeManagedPtr<ConnectionCallback> connStreamsCb,
       std::shared_ptr<const fizz::server::FizzServerContext> ctx,
       std::unique_ptr<CryptoFactory> cryptoFactory = nullptr,
       bool useConnectionEndWithErrorCallback = false);
@@ -92,8 +92,8 @@ class QuicServerTransport
   QuicServerTransport(
       folly::EventBase* evb,
       std::unique_ptr<QuicAsyncUDPSocketWrapper> sock,
-      ConnectionSetupCallback* connSetupCb,
-      ConnectionCallback* connStreamsCb,
+      folly::MaybeManagedPtr<ConnectionSetupCallback> connSetupCb,
+      folly::MaybeManagedPtr<ConnectionCallback> connStreamsCb,
       std::shared_ptr<const fizz::server::FizzServerContext> ctx,
       std::unique_ptr<CryptoFactory> cryptoFactory,
       PacketNum startingPacketNum);