StreamPool.c

 
#include <winpr/config.h>
#include <winpr/interlocked.h>
 
#include <winpr/crt.h>
#include <winpr/wlog.h>
 
#include <winpr/collections.h>
 
#include "../stream.h"
#include "../log.h"
#define XTAG WINPR_TAG("utils.streampool")
 
#if !defined(STREAMPOOL_SIZE_LIMIT)
#error "CMake must define STREAMPOOL_SIZE_LIMIT=<unsigned>"
#endif
static const size_t POOL_COMMON_LIMIT = STREAMPOOL_SIZE_LIMIT;
 
struct s_StreamPoolEntry
{
#if defined(WITH_DEBUG_STREAMPOOL)
  char** msg;
  size_t lines;
#endif
  wStream* s;
};
 
struct s_wStreamPool
{
  size_t aSize;
  size_t aCapacity;
  struct s_StreamPoolEntry* aArray;
 
  size_t uSize;
  size_t uCapacity;
  struct s_StreamPoolEntry* uArray;
 
  CRITICAL_SECTION lock;
  BOOL synchronized;
  size_t defaultSize;
  wLog* log;
};
 
static void discard_entry(struct s_StreamPoolEntry* entry, BOOL discardStream)
{
  if (!entry)
    return;
 
#if defined(WITH_DEBUG_STREAMPOOL)
  free((void*)entry->msg);
#endif
 
  if (discardStream && entry->s)
    Stream_Free(entry->s, entry->s->isAllocatedStream);
 
  const struct s_StreamPoolEntry empty = WINPR_C_ARRAY_INIT;
  *entry = empty;
}
 
static struct s_StreamPoolEntry add_entry(wStream* s)
{
  struct s_StreamPoolEntry entry = WINPR_C_ARRAY_INIT;
 
#if defined(WITH_DEBUG_STREAMPOOL)
  void* stack = winpr_backtrace(20);
  if (stack)
    entry.msg = winpr_backtrace_symbols(stack, &entry.lines);
  winpr_backtrace_free(stack);
#endif
 
  entry.s = s;
  return entry;
}
 
static inline void StreamPool_Lock(wStreamPool* pool)
{
  WINPR_ASSERT(pool);
  if (pool->synchronized)
    EnterCriticalSection(&pool->lock);
}
 
static inline void StreamPool_Unlock(wStreamPool* pool)
{
  WINPR_ASSERT(pool);
  if (pool->synchronized)
    LeaveCriticalSection(&pool->lock);
}
 
static BOOL StreamPool_ShrinkToCommonLimit(wStream* s)
{
  if (Stream_Capacity(s) <= POOL_COMMON_LIMIT)
    return TRUE;
 
  return Stream_ResizeToCapacity(s, POOL_COMMON_LIMIT);
}
 
static BOOL StreamPool_EnsureCapacity(wStreamPool* pool, size_t count, BOOL usedOrAvailable)
{
  WINPR_ASSERT(pool);
 
  size_t* cap = (usedOrAvailable) ? &pool->uCapacity : &pool->aCapacity;
  size_t* size = (usedOrAvailable) ? &pool->uSize : &pool->aSize;
  struct s_StreamPoolEntry** array = (usedOrAvailable) ? &pool->uArray : &pool->aArray;
 
  size_t new_cap = 0;
  if (*cap == 0)
    new_cap = *size + count;
  else if (*size + count > *cap)
    new_cap = (*size + count + 2) / 2 * 3;
  else if ((*size + count) < *cap / 3)
    new_cap = *cap / 2;
 
  if (new_cap > 0)
  {
    struct s_StreamPoolEntry* new_arr = nullptr;
 
    if (*cap < *size + count)
      *cap += count;
 
    new_arr =
        (struct s_StreamPoolEntry*)realloc(*array, sizeof(struct s_StreamPoolEntry) * new_cap);
    if (!new_arr)
      return FALSE;
    *cap = new_cap;
    *array = new_arr;
  }
  return TRUE;
}
 
static void StreamPool_ShiftUsed(wStreamPool* pool, size_t index)
{
  WINPR_ASSERT(pool);
 
  const size_t pcount = 1;
  const size_t off = index + pcount;
  if (pool->uSize >= off)
  {
    for (size_t x = 0; x < pcount; x++)
    {
      struct s_StreamPoolEntry* cur = &pool->uArray[index + x];
      discard_entry(cur, FALSE);
    }
    MoveMemory(&pool->uArray[index], &pool->uArray[index + pcount],
               (pool->uSize - index - pcount) * sizeof(struct s_StreamPoolEntry));
    pool->uSize -= pcount;
  }
}
 
static void StreamPool_AddUsed(wStreamPool* pool, wStream* s)
{
  StreamPool_EnsureCapacity(pool, 1, TRUE);
  pool->uArray[pool->uSize] = add_entry(s);
  pool->uSize++;
}
 
static void StreamPool_RemoveUsed(wStreamPool* pool, wStream* s)
{
  WINPR_ASSERT(pool);
  for (size_t index = 0; index < pool->uSize; index++)
  {
    struct s_StreamPoolEntry* cur = &pool->uArray[index];
    if (cur->s == s)
    {
      StreamPool_ShiftUsed(pool, index);
      break;
    }
  }
}
 
static void StreamPool_ShiftAvailable(wStreamPool* pool, size_t index)
{
  WINPR_ASSERT(pool);
 
  const size_t pcount = 1;
  const size_t off = index + pcount;
  if (pool->aSize >= off)
  {
    for (size_t x = 0; x < pcount; x++)
    {
      struct s_StreamPoolEntry* cur = &pool->aArray[index + x];
      discard_entry(cur, FALSE);
    }
 
    MoveMemory(&pool->aArray[index], &pool->aArray[index + pcount],
               (pool->aSize - index - pcount) * sizeof(struct s_StreamPoolEntry));
    pool->aSize -= pcount;
  }
}
 
wStream* StreamPool_Take(wStreamPool* pool, size_t size)
{
  BOOL found = FALSE;
  size_t foundIndex = 0;
  wStream* s = nullptr;
 
  StreamPool_Lock(pool);
 
  if (size == 0)
    size = pool->defaultSize;
 
  for (size_t index = 0; index < pool->aSize; index++)
  {
    struct s_StreamPoolEntry* cur = &pool->aArray[index];
    s = cur->s;
 
    if (Stream_Capacity(s) >= size)
    {
      found = TRUE;
      foundIndex = index;
      break;
    }
  }
 
  if (!found)
  {
    s = Stream_New(nullptr, size);
    if (!s)
      goto out_fail;
  }
  else if (s)
  {
    Stream_ResetPosition(s);
    if (!Stream_SetLength(s, Stream_Capacity(s)))
      goto out_fail;
    StreamPool_ShiftAvailable(pool, foundIndex);
  }
 
  if (s)
  {
    s->pool = pool;
    s->count = 1;
    StreamPool_AddUsed(pool, s);
  }
 
out_fail:
  StreamPool_Unlock(pool);
 
  return s;
}
 
static void StreamPool_Remove(wStreamPool* pool, wStream* s)
{
  StreamPool_EnsureCapacity(pool, 1, FALSE);
  Stream_EnsureValidity(s);
  StreamPool_ShrinkToCommonLimit(s);
  for (size_t x = 0; x < pool->aSize; x++)
  {
    wStream* cs = pool->aArray[x].s;
    if (cs == s)
      return;
  }
  pool->aArray[(pool->aSize)++] = add_entry(s);
  StreamPool_RemoveUsed(pool, s);
}
 
static void StreamPool_ReleaseOrReturn(wStreamPool* pool, wStream* s)
{
  StreamPool_Lock(pool);
  StreamPool_Remove(pool, s);
  StreamPool_Unlock(pool);
}
 
void StreamPool_Return(wStreamPool* pool, wStream* s)
{
  WINPR_ASSERT(pool);
  if (!s)
    return;
 
  StreamPool_Lock(pool);
  StreamPool_Remove(pool, s);
  StreamPool_Unlock(pool);
}
 
void Stream_AddRef(wStream* s)
{
  WINPR_ASSERT(s);
  InterlockedIncrement(&s->count);
}
 
void Stream_Release(wStream* s)
{
  WINPR_ASSERT(s);
 
  const LONG val = InterlockedDecrement(&s->count);
  if (val <= 0)
  {
    if (s->pool)
      StreamPool_ReleaseOrReturn(s->pool, s);
    else
      Stream_Free(s, TRUE);
  }
}
 
wStream* StreamPool_Find(wStreamPool* pool, const BYTE* ptr)
{
  wStream* s = nullptr;
 
  StreamPool_Lock(pool);
 
  for (size_t index = 0; index < pool->uSize; index++)
  {
    struct s_StreamPoolEntry* cur = &pool->uArray[index];
 
    if ((ptr >= Stream_Buffer(cur->s)) &&
        (ptr < (Stream_Buffer(cur->s) + Stream_Capacity(cur->s))))
    {
      s = cur->s;
      break;
    }
  }
 
  StreamPool_Unlock(pool);
 
  return s;
}
 
void StreamPool_Clear(wStreamPool* pool)
{
  StreamPool_Lock(pool);
 
  for (size_t x = 0; x < pool->aSize; x++)
  {
    struct s_StreamPoolEntry* cur = &pool->aArray[x];
    discard_entry(cur, TRUE);
  }
  pool->aSize = 0;
 
  if (pool->uSize > 0)
  {
    WLog_Print(pool->log, WLOG_WARN,
               "Clearing StreamPool, but there are %" PRIuz " streams currently in use",
               pool->uSize);
    for (size_t x = 0; x < pool->uSize; x++)
    {
      struct s_StreamPoolEntry* cur = &pool->uArray[x];
      discard_entry(cur, TRUE);
    }
    pool->uSize = 0;
  }
 
  StreamPool_Unlock(pool);
}
 
size_t StreamPool_UsedCount(wStreamPool* pool)
{
  StreamPool_Lock(pool);
  size_t usize = pool->uSize;
  StreamPool_Unlock(pool);
  return usize;
}
 
wStreamPool* StreamPool_New(BOOL synchronized, size_t defaultSize)
{
  wStreamPool* pool = calloc(1, sizeof(wStreamPool));
 
  if (!pool)
    return nullptr;
 
  pool->log = WLog_Create(XTAG, WLog_GetRoot());
  if (!pool->log)
    goto fail;
 
  pool->synchronized = synchronized;
  pool->defaultSize = defaultSize;
 
  if (!StreamPool_EnsureCapacity(pool, 32, FALSE))
    goto fail;
  if (!StreamPool_EnsureCapacity(pool, 32, TRUE))
    goto fail;
 
  if (!InitializeCriticalSectionAndSpinCount(&pool->lock, 4000))
    goto fail;
 
  return pool;
fail:
  WINPR_PRAGMA_DIAG_PUSH
  WINPR_PRAGMA_DIAG_IGNORED_MISMATCHED_DEALLOC
  StreamPool_Free(pool);
  WINPR_PRAGMA_DIAG_POP
  return nullptr;
}
 
void StreamPool_Free(wStreamPool* pool)
{
  if (!pool)
    return;
 
  StreamPool_Clear(pool);
 
  DeleteCriticalSection(&pool->lock);
 
  free(pool->aArray);
  free(pool->uArray);
 
  WLog_Discard(pool->log);
  free(pool);
}
 
char* StreamPool_GetStatistics(wStreamPool* pool, char* buffer, size_t size)
{
  WINPR_ASSERT(pool);
 
  if (!buffer || (size < 1))
    return nullptr;
 
  size_t used = 0;
  int offset = _snprintf(buffer, size - 1,
                         "aSize    =%" PRIuz ", uSize    =%" PRIuz ", aCapacity=%" PRIuz
                         ", uCapacity=%" PRIuz,
                         pool->aSize, pool->uSize, pool->aCapacity, pool->uCapacity);
  if ((offset > 0) && ((size_t)offset < size))
    used += (size_t)offset;
 
#if defined(WITH_DEBUG_STREAMPOOL)
  StreamPool_Lock(pool);
 
  offset = _snprintf(&buffer[used], size - 1 - used, "\n-- dump used array take locations --\n");
  if ((offset > 0) && ((size_t)offset < size - used))
    used += (size_t)offset;
  for (size_t x = 0; x < pool->uSize; x++)
  {
    const struct s_StreamPoolEntry* cur = &pool->uArray[x];
    WINPR_ASSERT(cur->msg || (cur->lines == 0));
 
    for (size_t y = 0; y < cur->lines; y++)
    {
      offset = _snprintf(&buffer[used], size - 1 - used, "[%" PRIuz " | %" PRIuz "]: %s\n", x,
                         y, cur->msg[y]);
      if ((offset > 0) && ((size_t)offset < size - used))
        used += (size_t)offset;
    }
  }
 
  offset = _snprintf(&buffer[used], size - 1 - used, "\n-- statistics called from --\n");
  if ((offset > 0) && ((size_t)offset < size - used))
    used += (size_t)offset;
 
  struct s_StreamPoolEntry entry = WINPR_C_ARRAY_INIT;
  void* stack = winpr_backtrace(20);
  if (stack)
    entry.msg = winpr_backtrace_symbols(stack, &entry.lines);
  winpr_backtrace_free(stack);
 
  for (size_t x = 0; x < entry.lines; x++)
  {
    const char* msg = entry.msg[x];
    offset = _snprintf(&buffer[used], size - 1 - used, "[%" PRIuz "]: %s\n", x, msg);
    if ((offset > 0) && ((size_t)offset < size - used))
      used += (size_t)offset;
  }
  free((void*)entry.msg);
  StreamPool_Unlock(pool);
#endif
  buffer[used] = '\0';
  return buffer;
}
 
BOOL StreamPool_WaitForReturn(wStreamPool* pool, UINT32 timeoutMS)
{
  /* HACK: We disconnected the transport above, now wait without a read or write lock until all
   * streams in use have been returned to the pool. */
  while (timeoutMS > 0)
  {
    const size_t used = StreamPool_UsedCount(pool);
    if (used == 0)
      return TRUE;
    WLog_Print(pool->log, WLOG_DEBUG, "%" PRIuz " streams still in use, sleeping...", used);
 
    char buffer[4096] = WINPR_C_ARRAY_INIT;
    StreamPool_GetStatistics(pool, buffer, sizeof(buffer));
    WLog_Print(pool->log, WLOG_TRACE, "Pool statistics: %s", buffer);
 
    UINT32 diff = 10;
    if (timeoutMS != INFINITE)
    {
      diff = timeoutMS > 10 ? 10 : timeoutMS;
      timeoutMS -= diff;
    }
    Sleep(diff);
  }
 
  return FALSE;
}