@@ -48,13 +48,14 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
4848 int ret = critnib_insert (hTracker -> map , (uintptr_t )ptr , value , 0 );
4949
5050 if (ret == 0 ) {
51- LOG_DEBUG ("memory region is added, tracker=%p, ptr=%p, size=%zu" ,
52- (void * )hTracker , ptr , size );
51+ LOG_DEBUG (
52+ "memory region is added, tracker=%p, ptr=%p, pool=%p, size=%zu" ,
53+ (void * )hTracker , ptr , (void * )pool , size );
5354 return UMF_RESULT_SUCCESS ;
5455 }
5556
56- LOG_ERR ("failed to insert tracker value, ret=%d, ptr=%p, size=%zu" , ret ,
57- ptr , size );
57+ LOG_ERR ("failed to insert tracker value, ret=%d, ptr=%p, pool=%p, size=%zu" ,
58+ ret , ptr , ( void * ) pool , size );
5859
5960 umf_ba_free (hTracker -> tracker_allocator , value );
6061
@@ -166,11 +167,35 @@ static umf_result_t trackingAlloc(void *hProvider, size_t size,
166167 return ret ;
167168 }
168169
169- umf_result_t ret2 = umfMemoryTrackerAdd (p -> hTracker , p -> pool , * ptr , size );
170- if (ret2 != UMF_RESULT_SUCCESS ) {
171- LOG_ERR ("failed to add allocated region to the tracker, ptr = %p, size "
170+ // check if the allocation was already added to the tracker
171+ // (in case of using ProxyLib)
172+ tracker_value_t * value =
173+ (tracker_value_t * )critnib_get (p -> hTracker -> map , * (uintptr_t * )ptr );
174+ if (value ) {
175+ assert (value -> pool != p -> pool );
176+
177+ LOG_DEBUG ("ptr already exists in the tracker (added by Proxy Lib) - "
178+ "updating value, ptr=%p, size=%zu, old pool: %p, new pool %p" ,
179+ * ptr , size , (void * )value -> pool , (void * )p -> pool );
180+
181+ // the allocation was made by the ProxyLib so we only update the tracker
182+ value -> pool = p -> pool ;
183+ int crit_ret = critnib_insert (p -> hTracker -> map , * (uintptr_t * )ptr ,
184+ value , 1 /* update */ );
185+
186+ // this cannot fail since we know the element exists and there is
187+ // nothing to allocate
188+ assert (crit_ret == 0 );
189+ (void )crit_ret ;
190+ } else {
191+ umf_result_t ret2 =
192+ umfMemoryTrackerAdd (p -> hTracker , p -> pool , * ptr , size );
193+ if (ret2 != UMF_RESULT_SUCCESS ) {
194+ LOG_ERR (
195+ "failed to add allocated region to the tracker, ptr = %p, size "
172196 "= %zu, ret = %d" ,
173197 * ptr , size , ret2 );
198+ }
174199 }
175200
176201 return ret ;
0 commit comments