src/liblzma/common/memory_limitter.c

   1 ///////////////////////////////////////////////////////////////////////////////
   2 //
   3 /// \file       memory_limitter.c
   4 /// \brief      Limitting memory usage
   5 //
   6 //  Copyright (C) 2007 Lasse Collin
   7 //
   8 //  This library is free software; you can redistribute it and/or
   9 //  modify it under the terms of the GNU Lesser General Public
  10 //  License as published by the Free Software Foundation; either
  11 //  version 2.1 of the License, or (at your option) any later version.
  12 //
  13 //  This library is distributed in the hope that it will be useful,
  14 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16 //  Lesser General Public License for more details.
  17 //
  18 ///////////////////////////////////////////////////////////////////////////////
  19
  20 #include "common.h"
  21
  22
  23 /// Rounds an unsigned integer upwards to the next multiple.
  24 #define my_ceil(num, multiple) \
  25         ((num) + (((multiple) - ((num) % (multiple))) % (multiple)))
  26
  27
  28 /// Rounds upwards to the next multiple of 2 * sizeof(void*).
  29 /// malloc() tends to align allocations this way.
  30 #define malloc_ceil(num) my_ceil(num, 2 * sizeof(void *))
  31
  32
  33 typedef struct lzma_memlimit_list_s lzma_memlimit_list;
  34 struct lzma_memlimit_list_s {
  35         lzma_memlimit_list *next;
  36         void *ptr;
  37         size_t size;
  38 };
  39
  40
  41 struct lzma_memlimit_s {
  42         size_t used;
  43         size_t limit;
  44         lzma_memlimit_list *list;
  45 };
  46
  47
  48 extern LZMA_API lzma_memlimit *
  49 lzma_memlimit_create(size_t limit)
  50 {
  51         if (limit < sizeof(lzma_memlimit))
  52                 return NULL;
  53
  54         lzma_memlimit *mem = malloc(sizeof(lzma_memlimit));
  55
  56         if (mem != NULL) {
  57                 mem->used = sizeof(lzma_memlimit);
  58                 mem->limit = limit;
  59                 mem->list = NULL;
  60         }
  61
  62         return mem;
  63 }
  64
  65
  66 extern LZMA_API void
  67 lzma_memlimit_set(lzma_memlimit *mem, size_t limit)
  68 {
  69         mem->limit = limit;
  70         return;
  71 }
  72
  73
  74 extern LZMA_API size_t
  75 lzma_memlimit_get(const lzma_memlimit *mem)
  76 {
  77         return mem->limit;
  78 }
  79
  80
  81 extern LZMA_API size_t
  82 lzma_memlimit_used(const lzma_memlimit *mem)
  83 {
  84         return mem->used;
  85 }
  86
  87
  88 extern LZMA_API size_t
  89 lzma_memlimit_count(const lzma_memlimit *mem)
  90 {
  91         // This is slow; we could have a counter in lzma_memlimit
  92         // for fast version. I expect the primary use of this
  93         // function to be limited to easy checking of memory leaks,
  94         // in which this implementation is just fine.
  95         size_t count = 0;
  96         const lzma_memlimit_list *record = mem->list;
  97
  98         while (record != NULL) {
  99                 ++count;
 100                 record = record->next;
 101         }
 102
 103         return count;
 104 }
 105
 106
 107 extern LZMA_API void
 108 lzma_memlimit_end(lzma_memlimit *mem, lzma_bool free_allocated)
 109 {
 110         if (mem == NULL)
 111                 return;
 112
 113         lzma_memlimit_list *record = mem->list;
 114         while (record != NULL) {
 115                 if (free_allocated)
 116                         free(record->ptr);
 117
 118                 lzma_memlimit_list *tmp = record;
 119                 record = record->next;
 120                 free(tmp);
 121         }
 122
 123         free(mem);
 124
 125         return;
 126 }
 127
 128
 129 extern LZMA_API void *
 130 lzma_memlimit_alloc(lzma_memlimit *mem, size_t nmemb, size_t size)
 131 {
 132         // While liblzma always sets nmemb to one, do this multiplication
 133         // to make these functions usable e.g. with zlib and libbzip2.
 134         // Making sure that this doesn't overflow is up to the application.
 135         size *= nmemb;
 136
 137         // Some malloc() implementations return NULL on malloc(0). We like
 138         // to get a non-NULL value.
 139         if (size == 0)
 140                 size = 1;
 141
 142         // Calculate how much memory we are going to allocate in reality.
 143         // TODO: We should add some rough estimate how much malloc() needs
 144         // for its internal structures.
 145         const size_t total_size = malloc_ceil(size)
 146                         + malloc_ceil(sizeof(lzma_memlimit_list));
 147
 148         // Integer overflow protection
 149         if (SIZE_MAX - size <= total_size)
 150                 return NULL;
 151
 152         if (mem->limit < mem->used || mem->limit - mem->used < total_size)
 153                 return NULL;
 154
 155         lzma_memlimit_list *record = malloc(sizeof(lzma_memlimit_list));
 156         void *ptr = malloc(size);
 157
 158         if (record == NULL || ptr == NULL) {
 159                 free(record);
 160                 free(ptr);
 161                 return NULL;
 162         }
 163
 164         // Add the new entry to the beginning of the list. This should be
 165         // more efficient when freeing memory, because usually it is
 166         // "last allocated, first freed".
 167         record->next = mem->list;
 168         record->ptr = ptr;
 169         record->size = total_size;
 170
 171         mem->list = record;
 172         mem->used += total_size;
 173
 174         return ptr;
 175 }
 176
 177
 178 extern LZMA_API void
 179 lzma_memlimit_detach(lzma_memlimit *mem, void *ptr)
 180 {
 181         if (ptr == NULL || mem->list == NULL)
 182                 return;
 183
 184         lzma_memlimit_list *record = mem->list;
 185         lzma_memlimit_list *prev = NULL;
 186
 187         while (record->ptr != ptr) {
 188                 prev = record;
 189                 record = record->next;
 190                 if (record == NULL)
 191                         return;
 192         }
 193
 194         if (prev != NULL)
 195                 prev->next = record->next;
 196         else
 197                 mem->list = record->next;
 198
 199         assert(mem->used >= record->size);
 200         mem->used -= record->size;
 201
 202         free(record);
 203
 204         return;
 205 }
 206
 207
 208 extern LZMA_API void
 209 lzma_memlimit_free(lzma_memlimit *mem, void *ptr)
 210 {
 211         if (ptr == NULL)
 212                 return;
 213
 214         lzma_memlimit_detach(mem, ptr);
 215
 216         free(ptr);
 217
 218         return;
 219 }