Fix memory leak in UnauthenticatedSessionTable. (#30025)

UnauthenticatedSessionTable essentially assumed that non-heap pools were used
and was:

1) Never releasing its entries back to the pool.
2) Assuming that the pool would fill up and then its "reuse already allocated
   entry with zero refcount" code would kick in.

Since heap pools never fill up, this meant that every single
UnauthenticatedSession allocated was leaked.  And we had a helpful "release them
all on destruction" to cover up the leak at shutdown and prevent leak tools from
finding it.

This fix:

* Preserves existing behavior for non-heap pools.
* Switches to releasing UnauthenticatedSessions back to the pool in the heap
  case.

Author: bzbarsky-apple

src/transport/UnauthenticatedSessionTable.h

@@ -22,6 +22,7 @@
 #include <lib/support/CodeUtils.h>
 #include <lib/support/Pool.h>
 #include <messaging/ReliableMessageProtocolConfig.h>
+#include <system/SystemConfig.h>
 #include <system/TimeSource.h>
 #include <transport/PeerMessageCounter.h>
 #include <transport/Session.h>
@@ -34,8 +35,7 @@ namespace Transport {
  * @brief
  *   An UnauthenticatedSession stores the binding of TransportAddress, and message counters.
  */
-class UnauthenticatedSession : public Session,
-                               public ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>
+class UnauthenticatedSession : public Session, public ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>
 {
 public:
     enum class SessionRole
@@ -44,6 +44,7 @@ class UnauthenticatedSession : public Session,
         kResponder,
     };
 
+protected:
     UnauthenticatedSession(SessionRole sessionRole, NodeId ephemeralInitiatorNodeID, const ReliableMessageProtocolConfig & config) :
         mEphemeralInitiatorNodeId(ephemeralInitiatorNodeID), mSessionRole(sessionRole),
         mLastActivityTime(System::SystemClock().GetMonotonicTimestamp()),
@@ -52,6 +53,7 @@ class UnauthenticatedSession : public Session,
     {}
     ~UnauthenticatedSession() override { VerifyOrDie(GetReferenceCount() == 0); }
 
+public:
     UnauthenticatedSession(const UnauthenticatedSession &)             = delete;
     UnauthenticatedSession & operator=(const UnauthenticatedSession &) = delete;
     UnauthenticatedSession(UnauthenticatedSession &&)                  = delete;
@@ -68,8 +70,8 @@ class UnauthenticatedSession : public Session,
 
     Session::SessionType GetSessionType() const override { return Session::SessionType::kUnauthenticated; }
 
-    void Retain() override { ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>::Retain(); }
-    void Release() override { ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>::Release(); }
+    void Retain() override { ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>::Retain(); }
+    void Release() override { ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>::Release(); }
 
     bool IsActiveSession() const override { return true; }
 
@@ -132,6 +134,23 @@ class UnauthenticatedSession : public Session,
 
     PeerMessageCounter & GetPeerMessageCounter() { return mPeerMessageCounter; }
 
+    static void Release(UnauthenticatedSession * obj)
+    {
+        // When using heap pools, we need to make sure to release ourselves back to
+        // the pool.  When not using heap pools, we don't want the extra cost of the
+        // table pointer here, and the table itself handles entry reuse and cleanup
+        // as needed.
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        obj->ReleaseSelfToPool();
+#else
+        // Just do nothing.
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    }
+
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    virtual void ReleaseSelfToPool() = 0;
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
 private:
     const NodeId mEphemeralInitiatorNodeId;
     const SessionRole mSessionRole;
@@ -142,6 +161,35 @@ class UnauthenticatedSession : public Session,
     PeerMessageCounter mPeerMessageCounter;
 };
 
+template <size_t kMaxSessionCount>
+class UnauthenticatedSessionTable;
+
+namespace detail {
+
+template <size_t kMaxSessionCount>
+class UnauthenticatedSessionPoolEntry : public UnauthenticatedSession
+{
+public:
+    UnauthenticatedSessionPoolEntry(SessionRole sessionRole, NodeId ephemeralInitiatorNodeID,
+                                    const ReliableMessageProtocolConfig & config,
+                                    UnauthenticatedSessionTable<kMaxSessionCount> & sessionTable) :
+        UnauthenticatedSession(sessionRole, ephemeralInitiatorNodeID, config)
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        ,
+        mSessionTable(sessionTable)
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    {}
+
+private:
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    virtual void ReleaseSelfToPool();
+
+    UnauthenticatedSessionTable<kMaxSessionCount> & mSessionTable;
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+};
+
+} // namespace detail
+
 /*
  * @brief
  *   An table which manages UnauthenticatedSessions
@@ -153,7 +201,17 @@ template <size_t kMaxSessionCount>
 class UnauthenticatedSessionTable
 {
 public:
-    ~UnauthenticatedSessionTable() { mEntries.ReleaseAll(); }
+    ~UnauthenticatedSessionTable()
+    {
+#if !CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        // When not using heap pools, our entries never actually get released
+        // back to the pool (which lets us make the entries 4 bytes smaller by
+        // not storing a reference to the table in them) and we LRU reuse ones
+        // that have 0 refcount.  But we should release them all here, to ensure
+        // that we don't hit fatal asserts in our pool destructor.
+        mEntries.ReleaseAll();
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    }
 
     /**
      * Get a responder session with the given ephemeralInitiatorNodeID. If the session doesn't exist in the cache, allocate a new
@@ -203,6 +261,9 @@ class UnauthenticatedSessionTable
     }
 
 private:
+    using EntryType = detail::UnauthenticatedSessionPoolEntry<kMaxSessionCount>;
+    friend EntryType;
+
     /**
      * Allocates a new session out of the internal resource pool.
      *
@@ -213,17 +274,23 @@ class UnauthenticatedSessionTable
     CHIP_ERROR AllocEntry(UnauthenticatedSession::SessionRole sessionRole, NodeId ephemeralInitiatorNodeID,
                           const ReliableMessageProtocolConfig & config, UnauthenticatedSession *& entry)
     {
-        entry = mEntries.CreateObject(sessionRole, ephemeralInitiatorNodeID, config);
-        if (entry != nullptr)
+        auto entryToUse = mEntries.CreateObject(sessionRole, ephemeralInitiatorNodeID, config, *this);
+        if (entryToUse != nullptr)
+        {
+            entry = entryToUse;
             return CHIP_NO_ERROR;
+        }
 
-        entry = FindLeastRecentUsedEntry();
-        if (entry == nullptr)
+#if !CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        entryToUse = FindLeastRecentUsedEntry();
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        if (entryToUse == nullptr)
         {
             return CHIP_ERROR_NO_MEMORY;
         }
 
-        mEntries.ResetObject(entry, sessionRole, ephemeralInitiatorNodeID, config);
+        mEntries.ResetObject(entryToUse, sessionRole, ephemeralInitiatorNodeID, config, *this);
+        entry = entryToUse;
         return CHIP_NO_ERROR;
     }
 
@@ -242,12 +309,12 @@ class UnauthenticatedSessionTable
         return result;
     }
 
-    UnauthenticatedSession * FindLeastRecentUsedEntry()
+    EntryType * FindLeastRecentUsedEntry()
     {
-        UnauthenticatedSession * result     = nullptr;
+        EntryType * result                  = nullptr;
         System::Clock::Timestamp oldestTime = System::Clock::Timestamp(std::numeric_limits<System::Clock::Timestamp::rep>::max());
 
-        mEntries.ForEachActiveObject([&](UnauthenticatedSession * entry) {
+        mEntries.ForEachActiveObject([&](EntryType * entry) {
             if (entry->GetReferenceCount() == 0 && entry->GetLastActivityTime() < oldestTime)
             {
                 result     = entry;
@@ -259,8 +326,18 @@ class UnauthenticatedSessionTable
         return result;
     }
 
-    ObjectPool<UnauthenticatedSession, kMaxSessionCount> mEntries;
+    void ReleaseEntry(EntryType * entry) { mEntries.ReleaseObject(entry); }
+
+    ObjectPool<EntryType, kMaxSessionCount> mEntries;
 };
 
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+template <size_t kMaxSessionCount>
+void detail::UnauthenticatedSessionPoolEntry<kMaxSessionCount>::ReleaseSelfToPool()
+{
+    mSessionTable.ReleaseEntry(this);
+}
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
 } // namespace Transport
 } // namespace chip