Pool.c 4.56 KB
Newer Older
Ben Gamari's avatar
Ben Gamari committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
/* ---------------------------------------------------------------------------
 *
 * (c) The GHC Team, 2014-2015
 *
 * A pool of lazily allocated things
 *
 * --------------------------------------------------------------------------*/

#include "PosixSource.h"
#include "Rts.h"
#include "RtsUtils.h"
#include "Pool.h"

/* used to mark an entry as needing to be freed when released */
#define FLAG_SHOULD_FREE (1 << 0)

typedef struct PoolEntry_ {
    struct PoolEntry_ *next;
    void *thing;
    StgWord flags;
} PoolEntry;

struct Pool_ {
    /* the maximum number of allocated resources in the pool */
    nat max_size;
    /* the number of allocated resources to keep in the pool when idle */
    nat desired_size;
    /* how many things are currently allocated? (sum of lengths of available and
     * taken lists) */
    nat current_size;
#ifdef THREADED_RTS
    /* signaled when a thing is released */
    Condition cond;
#endif
    alloc_thing_fn alloc_fn;
    free_thing_fn free_fn;

    PoolEntry *available;
    PoolEntry *taken;
#ifdef THREADED_RTS
    /* protects entire data structure */
    Mutex mutex;
#endif
};

Pool *poolInit(nat max_size, nat desired_size,
               alloc_thing_fn alloc_fn, free_thing_fn free_fn) {
    Pool *pool = stgMallocBytes(sizeof(Pool), "pool_init");
    pool->max_size = max_size == 0 ? (nat) -1 : max_size;
    pool->desired_size = desired_size;
    pool->current_size = 0;
    pool->alloc_fn = alloc_fn;
    pool->free_fn = free_fn;
    pool->available = NULL;
    pool->taken = NULL;
#ifdef THREADED_RTS
    initMutex(&pool->mutex);
    initCondition(&pool->cond);
#endif
    return pool;
}

int poolFree(Pool *pool) {
    if (pool->taken != NULL)
        return 1;

    poolSetMaxSize(pool, 0);
#ifdef THREADED_RTS
    closeCondition(&pool->cond);
    closeMutex(&pool->mutex);
#endif
    free(pool);
    return 0;
}

/* free available entries such that current_size <= size */
static void free_available(Pool *pool, nat size) {
    while (pool->current_size > size && pool->available != NULL) {
        PoolEntry *ent = pool->available;
        pool->free_fn(ent->thing);
        pool->available = ent->next;
        free(ent);
        pool->current_size--;
    }
}

void poolSetDesiredSize(Pool *pool, nat size) {
    ACQUIRE_LOCK(&pool->mutex);
    pool->desired_size = size;
    free_available(pool, size);
    RELEASE_LOCK(&pool->mutex);
}

void poolSetMaxSize(Pool *pool, nat size) {
    ACQUIRE_LOCK(&pool->mutex);
    if (size == 0)
        size = (nat) -1;
    pool->max_size = size;
    if (pool->desired_size > pool->max_size) {
        pool->desired_size = size;
        free_available(pool, size);
    }
    RELEASE_LOCK(&pool->mutex);
}

nat poolGetMaxSize(Pool *pool) {
    return pool->max_size;
}

nat poolGetDesiredSize(Pool *pool) {
    return pool->desired_size;
}

void *poolTake(Pool *pool) {
    PoolEntry *ent = NULL;
    ACQUIRE_LOCK(&pool->mutex);
    while (ent == NULL) {
        if (pool->available != NULL) {
            ent = pool->available;
            pool->available = ent->next;
        } else if (pool->current_size < pool->max_size) {
            ent = stgMallocBytes(sizeof(PoolEntry), "pool_take");
            ent->flags = 0;
            ent->thing = pool->alloc_fn();
            pool->current_size++;
        } else {
#ifdef THREADED_RTS
            waitCondition(&pool->cond, &pool->mutex);
#else
            barf("Tried to take from an empty pool");
#endif
        }
    }

    ent->next = pool->taken;
    pool->taken = ent;
    RELEASE_LOCK(&pool->mutex);
    return ent->thing;
}

void poolRelease(Pool *pool, void *thing) {
    ACQUIRE_LOCK(&pool->mutex);
    PoolEntry **last = &pool->taken;
    PoolEntry *ent = pool->taken;
    while (ent != NULL) {
        if (ent->thing == thing) {
            *last = ent->next;
            if (pool->current_size > pool->desired_size
                || ent->flags & FLAG_SHOULD_FREE) {
                pool->free_fn(ent->thing);
                free(ent);
            } else {
                ent->next = pool->available;
                pool->available = ent;
#ifdef THREADED_RTS
                signalCondition(&pool->cond);
#endif
            }

            RELEASE_LOCK(&pool->mutex);
            return;
        }

        last = &ent->next;
        ent = ent->next;
    }

    barf("pool_release: trying to release resource which doesn't belong to pool.");
}

void poolFlush(Pool *pool) {
    ACQUIRE_LOCK(&pool->mutex);
    free_available(pool, 0);
    PoolEntry *ent = pool->taken;
    while (ent != NULL) {
        ent->flags |= FLAG_SHOULD_FREE;
        ent = ent->next;
    }
    RELEASE_LOCK(&pool->mutex);
}