Use ReferenceCounted for ExchangeContext (#3511)

* Use ReferenceCounted for ExchangeContext

Refactor only, no behavior changes.

* Restyled by clang-format

* Fix naming

Co-authored-by: Restyled.io <commits@restyled.io>

src/lib/core/ReferenceCounted.h

@@ -41,22 +41,22 @@ class DeleteDeletor
  * A reference counted object maintains a count of usages and when the usage
  * count drops to 0, it deletes itself.
  */
-template <class SUBCLASS, class DELETOR = DeleteDeletor<SUBCLASS>>
+template <class Subclass, class Deletor = DeleteDeletor<Subclass>, int kInitRefCount = 1>
 class ReferenceCounted
 {
 public:
     typedef uint32_t count_type;
 
     /** Adds one to the usage count of this class */
-    SUBCLASS * Retain()
+    Subclass * Retain()
     {
         if (mRefCount == std::numeric_limits<count_type>::max())
         {
             abort();
         }
         ++mRefCount;
 
-        return static_cast<SUBCLASS *>(this);
+        return static_cast<Subclass *>(this);
     }
 
     /** Release usage of this class */
@@ -69,15 +69,15 @@ class ReferenceCounted
 
         if (--mRefCount == 0)
         {
-            DELETOR::Release(static_cast<SUBCLASS *>(this));
+            Deletor::Release(static_cast<Subclass *>(this));
         }
     }
 
     /** Get the current reference counter value */
     count_type GetReferenceCount() const { return mRefCount; }
 
 private:
-    count_type mRefCount = 1;
+    count_type mRefCount = kInitRefCount;
 };
 
 } // namespace chip

src/messaging/ExchangeContext.cpp

@@ -82,12 +82,12 @@ CHIP_ERROR ExchangeContext::SendMessage(uint16_t protocolId, uint8_t msgType, Pa
     PayloadHeader payloadHeader;
 
     // Don't let method get called on a freed object.
-    VerifyOrDie(mExchangeMgr != nullptr && mRefCount != 0);
+    VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() > 0);
 
     // we hold the exchange context here in case the entity that
     // originally generated it tries to close it as a result of
     // an error arising below. at the end, we have to close it.
-    AddRef();
+    Retain();
 
     // Set the exchange ID for this header.
     payloadHeader.SetExchangeID(mExchangeId);
@@ -141,19 +141,6 @@ CHIP_ERROR ExchangeContext::SendMessage(uint16_t protocolId, uint8_t msgType, Pa
     return err;
 }
 
-/**
- *  Increment the reference counter for the exchange context by one.
- *
- */
-void ExchangeContext::AddRef()
-{
-    mRefCount++;
-#if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
-    ChipLogProgress(ExchangeManager, "ec id: %d [%04" PRIX16 "], refCount++: %d", (this - mExchangeMgr->ContextPool + 1),
-                    mExchangeId, mRefCount);
-#endif
-}
-
 void ExchangeContext::DoClose(bool clearRetransTable)
 {
     // Clear protocol callbacks
@@ -171,7 +158,7 @@ void ExchangeContext::DoClose(bool clearRetransTable)
  */
 void ExchangeContext::Close()
 {
-    VerifyOrDie(mExchangeMgr != nullptr && mRefCount != 0);
+    VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() > 0);
 
 #if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
     ChipLogProgress(ExchangeManager, "ec id: %d [%04" PRIX16 "], %s", (this - mExchangeMgr->ContextPool + 1), mExchangeId,
@@ -188,7 +175,7 @@ void ExchangeContext::Close()
  */
 void ExchangeContext::Abort()
 {
-    VerifyOrDie(mExchangeMgr != nullptr && mRefCount != 0);
+    VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() > 0);
 
 #if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
     ChipLogProgress(ExchangeManager, "ec id: %d [%04" PRIX16 "], %s", (this - mExchangeMgr->ContextPool + 1), mExchangeId,
@@ -199,46 +186,50 @@ void ExchangeContext::Abort()
     Release();
 }
 
-/**
- *  Release reference to this exchange context. If count is down
- *  to one then close the context, reset all protocol callbacks,
- *  and stop all timers.
- *
- */
-void ExchangeContext::Release()
+void ExchangeContext::Reset()
 {
-    VerifyOrDie(mExchangeMgr != nullptr && mRefCount != 0);
+    *this = ExchangeContext();
+}
+
+void ExchangeContext::Alloc(ExchangeManager * em, uint16_t ExchangeId, uint64_t PeerNodeId, bool Initiator, void * AppState)
+{
+    VerifyOrDie(mExchangeMgr == nullptr && GetReferenceCount() == 0);
+
+    Reset();
+    Retain();
+    mExchangeMgr = em;
+    em->IncrementContextsInUse();
+    mExchangeId = ExchangeId;
+    mPeerNodeId = PeerNodeId;
+    mFlags.Set(ExFlagValues::kFlagInitiator, Initiator);
+    mAppState = AppState;
 
-    if (mRefCount == 1)
-    {
-        // Ideally, in this scenario, the retransmit table should
-        // be clear of any outstanding messages for this context and
-        // the boolean parameter passed to DoClose() should not matter.
-        ExchangeManager * em = mExchangeMgr;
 #if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
-        uint16_t tmpid = mExchangeId;
+    ChipLogProgress(ExchangeManager, "ec++ id: %d, inUse: %d, addr: 0x%x", (this - em->ContextPool + 1), em->GetContextsInUse(),
+                    this);
 #endif
+    SYSTEM_STATS_INCREMENT(chip::System::Stats::kExchangeMgr_NumContexts);
+}
 
-        DoClose(false);
-        mRefCount    = 0;
-        mExchangeMgr = nullptr;
+void ExchangeContext::Free()
+{
+    VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() == 0);
 
-        em->DecrementContextsInUse();
+    // Ideally, in this scenario, the retransmit table should
+    // be clear of any outstanding messages for this context and
+    // the boolean parameter passed to DoClose() should not matter.
+    ExchangeManager * em = mExchangeMgr;
+
+    DoClose(false);
+    mExchangeMgr = nullptr;
+
+    em->DecrementContextsInUse();
 
 #if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
-        ChipLogProgress(ExchangeManager, "ec-- id: %d [%04" PRIX16 "], inUse: %d, addr: 0x%x", (this - em->ContextPool + 1), tmpid,
-                        em->GetContextsInUse(), this);
-#endif
-        SYSTEM_STATS_DECREMENT(chip::System::Stats::kExchangeMgr_NumContexts);
-    }
-    else
-    {
-        mRefCount--;
-#if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
-        ChipLogProgress(ExchangeManager, "ec id: %d [%04" PRIX16 "], refCount--: %d", (this - mExchangeMgr->ContextPool + 1),
-                        mExchangeId, mRefCount);
+    ChipLogProgress(ExchangeManager, "ec-- id: %d [%04" PRIX16 "], inUse: %d, addr: 0x%x", (this - em->ContextPool + 1),
+                    mExchangeId, em->GetContextsInUse(), this);
 #endif
-    }
+    SYSTEM_STATS_DECREMENT(chip::System::Stats::kExchangeMgr_NumContexts);
 }
 
 bool ExchangeContext::MatchExchange(const PacketHeader & packetHeader, const PayloadHeader & payloadHeader)
@@ -317,7 +308,7 @@ CHIP_ERROR ExchangeContext::HandleMessage(const PacketHeader & packetHeader, con
     // guard against Close() calls(decrementing the reference
     // count) by the protocol before the CHIP Exchange
     // layer has completed its work on the ExchangeContext.
-    AddRef();
+    Retain();
 
     protocolId  = payloadHeader.GetProtocolID();
     messageType = payloadHeader.GetMessageType();

src/messaging/ExchangeContext.h

@@ -23,6 +23,7 @@
 
 #pragma once
 
+#include <lib/core/ReferenceCounted.h>
 #include <support/BitFlags.h>
 #include <support/DLLUtil.h>
 #include <system/SystemTimer.h>
@@ -68,15 +69,22 @@ class DLL_EXPORT ExchangeContextDelegate
     virtual void OnResponseTimeout(ExchangeContext * ec) = 0;
 };
 
+class ExchangeContextDeletor
+{
+public:
+    static void Release(ExchangeContext * obj);
+};
+
 /**
  *  @brief
  *    This class represents an ongoing conversation (ExchangeContext) between two or more nodes.
  *    It defines methods for encoding and communicating CHIP messages within an ExchangeContext
  *    over various transport mechanisms, for example, TCP, UDP, or CHIP Reliable Messaging.
  */
-class DLL_EXPORT ExchangeContext
+class DLL_EXPORT ExchangeContext : public ReferenceCounted<ExchangeContext, ExchangeContextDeletor, 0>
 {
     friend class ExchangeManager;
+    friend class ExchangeContextDeletor;
 
 public:
     enum
@@ -205,11 +213,12 @@ class DLL_EXPORT ExchangeContext
      * can hold onto it while it's out of their direct control to make sure it isn't closed before everyone's ready.
      * A customized version of reference counting is used since there are some extra stuff to do within Release.
      */
-    void AddRef();
     void Close();
     void Abort();
-    void Release();
-    void SetRefCount(uint8_t value) { mRefCount = value; }
+
+    void Alloc(ExchangeManager * em, uint16_t ExchangeId, uint64_t PeerNodeId, bool Initiator, void * AppState);
+    void Free();
+    void Reset();
 
 private:
     enum class ExFlagValues : uint16_t
@@ -227,7 +236,6 @@ class DLL_EXPORT ExchangeContext
 
     uint64_t mPeerNodeId; // Node ID of peer node.
     uint16_t mExchangeId; // Assigned exchange ID.
-    uint8_t mRefCount;    // Reference counter of the current instance
 
     BitFlags<uint16_t, ExFlagValues> mFlags; // Internal state flags
 
@@ -256,4 +264,9 @@ class DLL_EXPORT ExchangeContext
     void DoClose(bool clearRetransTable);
 };
 
+inline void ExchangeContextDeletor::Release(ExchangeContext * obj)
+{
+    obj->Free();
+}
+
 } // namespace chip

src/messaging/ExchangeMgr.cpp

@@ -101,18 +101,7 @@ CHIP_ERROR ExchangeManager::Shutdown()
 
 ExchangeContext * ExchangeManager::NewContext(const NodeId & peerNodeId, void * appState)
 {
-    ExchangeContext * ec = AllocContext();
-
-    if (ec != nullptr)
-    {
-        ec->SetExchangeId(mNextExchangeId++);
-        ec->SetPeerNodeId(peerNodeId);
-        ec->SetAppState(appState);
-        ec->SetInitiator(true);
-        ChipLogProgress(ExchangeManager, "ec id: %d, AppState: 0x%x", (ec - ContextPool + 1), ec->GetAppState());
-    }
-
-    return ec;
+    return AllocContext(mNextExchangeId++, peerNodeId, true, appState);
 }
 
 ExchangeContext * ExchangeManager::FindContext(NodeId peerNodeId, void * appState, bool isInitiator)
@@ -121,7 +110,7 @@ ExchangeContext * ExchangeManager::FindContext(NodeId peerNodeId, void * appStat
 
     for (int i = 0; i < CHIP_CONFIG_MAX_EXCHANGE_CONTEXTS; i++, ec++)
     {
-        if (ec->GetExchangeMgr() != nullptr && ec->GetPeerNodeId() == peerNodeId && ec->GetAppState() == appState &&
+        if (ec->GetReferenceCount() > 0 && ec->GetPeerNodeId() == peerNodeId && ec->GetAppState() == appState &&
             ec->IsInitiator() == isInitiator)
             return ec;
     }
@@ -156,26 +145,17 @@ void ExchangeManager::OnReceiveError(CHIP_ERROR error, const Transport::PeerAddr
     ChipLogError(ExchangeManager, "Accept FAILED, err = %s", ErrorStr(error));
 }
 
-ExchangeContext * ExchangeManager::AllocContext()
+ExchangeContext * ExchangeManager::AllocContext(uint16_t ExchangeId, uint64_t PeerNodeId, bool Initiator, void * AppState)
 {
     ExchangeContext * ec = ContextPool;
 
     CHIP_FAULT_INJECT(FaultInjection::kFault_AllocExchangeContext, return nullptr);
 
     for (int i = 0; i < CHIP_CONFIG_MAX_EXCHANGE_CONTEXTS; i++, ec++)
     {
-        if (ec->GetExchangeMgr() == nullptr)
+        if (ec->GetReferenceCount() == 0)
         {
-            *ec = ExchangeContext();
-            ec->SetExchangeMgr(this);
-            ec->SetRefCount(1);
-            mContextsInUse++;
-
-#if defined(CHIP_EXCHANGE_CONTEXT_DETAIL_LOGGING)
-            ChipLogProgress(ExchangeManager, "ec++ id: %d, inUse: %d, addr: 0x%x", (ec - ContextPool + 1), mContextsInUse, ec);
-#endif
-            SYSTEM_STATS_INCREMENT(chip::System::Stats::kExchangeMgr_NumContexts);
-
+            ec->Alloc(this, ExchangeId, PeerNodeId, Initiator, AppState);
             return ec;
         }
     }
@@ -195,7 +175,7 @@ void ExchangeManager::DispatchMessage(const PacketHeader & packetHeader, const P
     ec = ContextPool;
     for (int i = 0; i < CHIP_CONFIG_MAX_EXCHANGE_CONTEXTS; i++, ec++)
     {
-        if (ec->GetExchangeMgr() != nullptr && ec->MatchExchange(packetHeader, payloadHeader))
+        if (ec->GetReferenceCount() > 0 && ec->MatchExchange(packetHeader, payloadHeader))
         {
             // Matched ExchangeContext; send to message handler.
             ec->HandleMessage(packetHeader, payloadHeader, msgBuf);
@@ -243,14 +223,9 @@ void ExchangeManager::DispatchMessage(const PacketHeader & packetHeader, const P
     {
         ExchangeContext::MessageReceiveFunct umhandler = nullptr;
 
-        ec = AllocContext();
+        ec = AllocContext(payloadHeader.GetExchangeID(), packetHeader.GetSourceNodeId().Value(), false, matchingUMH->AppState);
         VerifyOrExit(ec != nullptr, err = CHIP_ERROR_NO_MEMORY);
 
-        ec->SetExchangeId(payloadHeader.GetExchangeID());
-        ec->SetPeerNodeId(packetHeader.GetSourceNodeId().Value());
-        ec->SetInitiator(false);
-        ec->SetAppState(matchingUMH->AppState);
-
         umhandler = matchingUMH->Handler;
 
         ChipLogProgress(ExchangeManager, "ec id: %d, AppState: 0x%x", (ec - ContextPool + 1), ec->GetAppState());
@@ -329,6 +304,11 @@ void ExchangeManager::OnMessageReceived(const PacketHeader & packetHeader, const
     DispatchMessage(packetHeader, payloadHeader, msgBuf);
 }
 
+void ExchangeManager::IncrementContextsInUse()
+{
+    mContextsInUse++;
+}
+
 void ExchangeManager::DecrementContextsInUse()
 {
     if (mContextsInUse >= 1)

src/messaging/ExchangeMgr.h

@@ -160,9 +160,7 @@ class DLL_EXPORT ExchangeManager : public SecureSessionMgrDelegate
      */
     CHIP_ERROR UnregisterUnsolicitedMessageHandler(uint32_t protocolId, uint8_t msgType);
 
-    /**
-     *  Decrement current context in use by 1.
-     */
+    void IncrementContextsInUse();
     void DecrementContextsInUse();
 
     SecureSessionMgrBase * GetSessionMgr() const { return mSessionMgr; }
@@ -194,7 +192,7 @@ class DLL_EXPORT ExchangeManager : public SecureSessionMgrDelegate
     UnsolicitedMessageHandler UMHandlerPool[CHIP_CONFIG_MAX_UNSOLICITED_MESSAGE_HANDLERS];
     void (*OnExchangeContextChanged)(size_t numContextsInUse);
 
-    ExchangeContext * AllocContext();
+    ExchangeContext * AllocContext(uint16_t ExchangeId, uint64_t PeerNodeId, bool Initiator, void * AppState);
 
     void DispatchMessage(const PacketHeader & packetHeader, const PayloadHeader & payloadHeader, System::PacketBuffer * msgBuf);