PRD: Add NULL checks for memory in FlyWeight

Change-Id: If655cf403ace276f9d9a0d71ed016a5316148204
Reviewed-on: http://gfw160.austin.ibm.com:8080/gerrit/3617
Tested-by: Jenkins Server
Reviewed-by: Christopher T. Phan <cphan@us.ibm.com>
Reviewed-by: Zane Shelley <zshelle@us.ibm.com>
Reviewed-by: A. Patrick Williams III <iawillia@us.ibm.com>
Reviewed-on: http://gfw160.austin.ibm.com:8080/gerrit/3829

2 files changed, 19 insertions, 3 deletions

diff --git a/src/usr/diag/prdf/common/util/prdfFlyWeight.C b/src/usr/diag/prdf/common/util/prdfFlyWeight.C
index 151e49ae6..05df37a00 100755
--- a/src/usr/diag/prdf/common/util/prdfFlyWeight.C
+++ b/src/usr/diag/prdf/common/util/prdfFlyWeight.C
@@ -36,6 +36,7 @@
 #include <tracinterface.H>
 #include <prdfHeapBucketSize.H>
 #include <prdfTrace.H>
+#include <prdfAssert.h>
 
 #undef prdfFlyWeight_C
 //----------------------------------------------------------------------
@@ -114,7 +115,8 @@ T & FlyWeight<T,S>::get(const T & key)
       current_array = new T[RoundBucketSize<T,S>::value];
       if(current_array == NULL)   // dg00a
       {                           // dg00a
-        PRDF_TRAC("PRDF Could not get requested memory");  // dg00a
+        PRDF_ERR("PRDF Could not get requested memory");
+        PRDF_ASSERT( NULL != current_array );
       }                           // dg00a
       // if the heap of array ptrs is full(or non-existant) then increase capacity by S dg00a
       if(iv_heap.size() == iv_heap.capacity())     // dg00a
diff --git a/src/usr/diag/prdf/common/util/prdfFlyWeightS.C b/src/usr/diag/prdf/common/util/prdfFlyWeightS.C
index 8a8a73789..325fa729b 100755
--- a/src/usr/diag/prdf/common/util/prdfFlyWeightS.C
+++ b/src/usr/diag/prdf/common/util/prdfFlyWeightS.C
@@ -24,6 +24,8 @@
 #include <prdfFlyWeightS.H>
 #include <algorithm>
 #include <prdfHeapBucketSize.H>
+#include <prdfTrace.H>
+#include <prdfAssert.h>
 
 namespace PRDF
 {
@@ -201,7 +203,13 @@ T * FlyWeightS<T,S>::getHeap()
 
     if (NULL == iv_nextPos)
     {
-        iv_heaps.push_back(iv_nextPos = new T[RoundBucketSize<T,S>::value]);
+        iv_nextPos = new T[RoundBucketSize<T,S>::value];
+        if( NULL == iv_nextPos )
+        {
+            PRDF_ERR("Can not allocate memory on heap");
+            PRDF_ASSERT( NULL != iv_nextPos );
+        }
+        iv_heaps.push_back(iv_nextPos);
     }
 
     T * l_rc = iv_nextPos;
@@ -209,7 +217,13 @@ T * FlyWeightS<T,S>::getHeap()
     iv_nextPos++;
     if ((static_cast<T*>(*(iv_heaps.end()-1)) + RoundBucketSize<T,S>::value) == iv_nextPos)
     {
-        iv_heaps.push_back(iv_nextPos = new T[RoundBucketSize<T,S>::value]);
+        iv_nextPos = new T[RoundBucketSize<T,S>::value];
+        if( NULL == iv_nextPos )
+        {
+            PRDF_ERR("Can not allocate memory on heap");
+            PRDF_ASSERT( NULL != iv_nextPos );
+        }
+        iv_heaps.push_back(iv_nextPos);
     }
 
     return l_rc;