From cf1d579462815dcfe3e2ebad92e18ed9ea449b68 Mon Sep 17 00:00:00 2001
From: Lorenzo Blasa <lblasa@fb.com>
Date: Wed, 29 Sep 2021 07:43:01 -0700
Subject: [PATCH] Manual disconnect on deallocation and do not clear message
 handler or message received

Summary:
Trigger a manual disconnect on deallocation. This was done automatically for us when the underlying socket gets released. But, this gives a bit more visibility and control onto exactly when this is going to take place.

Additionally, do not clear the message handler when a message is received.

It is not required as sendExpectResponse is one time called only used for certificate exchange. If this takes place again, a new handler will be set anyway.

Reviewed By: passy

Differential Revision: D31231828

fbshipit-source-id: 36ad13564a358b88d1618e94195fe05433d80993
---
 iOS/FlipperKit/FlipperWebSocket.h  | 18 ++++++------------
 iOS/FlipperKit/FlipperWebSocket.mm | 17 +++++++----------
 2 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/iOS/FlipperKit/FlipperWebSocket.h b/iOS/FlipperKit/FlipperWebSocket.h
index 119ca3064..8a661922f 100644
--- a/iOS/FlipperKit/FlipperWebSocket.h
+++ b/iOS/FlipperKit/FlipperWebSocket.h
@@ -28,15 +28,13 @@ class FlipperWebSocket : public FlipperSocket {
  public:
   FlipperWebSocket(
       FlipperConnectionEndpoint endpoint,
-      std::unique_ptr<FlipperSocketBasePayload> payload,
-      folly::EventBase* eventBase);
+      std::unique_ptr<FlipperSocketBasePayload> payload);
   FlipperWebSocket(
       FlipperConnectionEndpoint endpoint,
       std::unique_ptr<FlipperSocketBasePayload> payload,
-      folly::EventBase* eventBase,
       ConnectionContextStore* connectionContextStore);
 
-  virtual ~FlipperWebSocket() {}
+  virtual ~FlipperWebSocket();
 
   virtual void setEventHandler(SocketEventHandler eventHandler) override;
   virtual void setMessageHandler(SocketMessageHandler messageHandler) override;
@@ -55,13 +53,12 @@ class FlipperWebSocket : public FlipperSocket {
  private:
   FlipperConnectionEndpoint endpoint_;
   std::unique_ptr<FlipperSocketBasePayload> payload_;
-  folly::EventBase* eventBase_;
   ConnectionContextStore* connectionContextStore_;
 
-  FlipperPlatformWebSocket* socket_;
-
   SocketEventHandler eventHandler_;
   SocketMessageHandler messageHandler_;
+
+  FlipperPlatformWebSocket* socket_;
 };
 
 class FlipperWebSocketProvider : public FlipperSocketProvider {
@@ -72,7 +69,7 @@ class FlipperWebSocketProvider : public FlipperSocketProvider {
       std::unique_ptr<FlipperSocketBasePayload> payload,
       folly::EventBase* eventBase) override {
     return std::make_unique<FlipperWebSocket>(
-        std::move(endpoint), std::move(payload), eventBase);
+        std::move(endpoint), std::move(payload));
   }
   virtual std::unique_ptr<FlipperSocket> create(
       FlipperConnectionEndpoint endpoint,
@@ -80,10 +77,7 @@ class FlipperWebSocketProvider : public FlipperSocketProvider {
       folly::EventBase* eventBase,
       ConnectionContextStore* connectionContextStore) override {
     return std::make_unique<FlipperWebSocket>(
-        std::move(endpoint),
-        std::move(payload),
-        eventBase,
-        connectionContextStore);
+        std::move(endpoint), std::move(payload), connectionContextStore);
   }
 };
 
diff --git a/iOS/FlipperKit/FlipperWebSocket.mm b/iOS/FlipperKit/FlipperWebSocket.mm
index 784844521..d73f429b1 100644
--- a/iOS/FlipperKit/FlipperWebSocket.mm
+++ b/iOS/FlipperKit/FlipperWebSocket.mm
@@ -92,22 +92,21 @@ class WebSocketSerializer : public FlipperPayloadSerializer {
 
 FlipperWebSocket::FlipperWebSocket(
     FlipperConnectionEndpoint endpoint,
-    std::unique_ptr<FlipperSocketBasePayload> payload,
-    folly::EventBase* eventBase)
-    : endpoint_(std::move(endpoint)),
-      payload_(std::move(payload)),
-      eventBase_(eventBase) {}
+    std::unique_ptr<FlipperSocketBasePayload> payload)
+    : endpoint_(std::move(endpoint)), payload_(std::move(payload)) {}
 
 FlipperWebSocket::FlipperWebSocket(
     FlipperConnectionEndpoint endpoint,
     std::unique_ptr<FlipperSocketBasePayload> payload,
-    folly::EventBase* eventBase,
     ConnectionContextStore* connectionContextStore)
     : endpoint_(std::move(endpoint)),
       payload_(std::move(payload)),
-      eventBase_(eventBase),
       connectionContextStore_(connectionContextStore) {}
 
+FlipperWebSocket::~FlipperWebSocket() {
+  disconnect();
+}
+
 void FlipperWebSocket::setEventHandler(SocketEventHandler eventHandler) {
   eventHandler_ = std::move(eventHandler);
 }
@@ -169,8 +168,7 @@ bool FlipperWebSocket::connect(FlipperConnectionManager* manager) {
         promise.set_value(false);
       }
     }
-    eventBase_->runInEventBaseThread(
-        [eventHandler, event]() { eventHandler(event); });
+    eventHandler(event);
   };
   socket_.messageHandler = ^(const std::string& message) {
     this->messageHandler_(message);
@@ -240,7 +238,6 @@ void FlipperWebSocket::sendExpectResponse(
 
   [socket_ setMessageHandler:^(const std::string& msg) {
     completion(msg, false);
-    [socket_ setMessageHandler:NULL];
   }];
   NSError* error = NULL;
   [socket_ send:messageObjc error:&error];