FlashGraph-ng
A new frontier in large-scale graph analysis and data mining
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Pages
slab_allocator.h
1 #ifndef __SLAB_ALLOCATOR_H__
2 #define __SLAB_ALLOCATOR_H__
3 
4 /*
5  * Copyright 2014 Open Connectome Project (http://openconnecto.me)
6  * Written by Da Zheng (zhengda1936@gmail.com)
7  *
8  * This file is part of SAFSlib.
9  *
10  * Licensed under the Apache License, Version 2.0 (the "License");
11  * you may not use this file except in compliance with the License.
12  * You may obtain a copy of the License at
13  *
14  * http://www.apache.org/licenses/LICENSE-2.0
15  *
16  * Unless required by applicable law or agreed to in writing, software
17  * distributed under the License is distributed on an "AS IS" BASIS,
18  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19  * See the License for the specific language governing permissions and
20  * limitations under the License.
21  */
22 
23 #include <stdlib.h>
24 #include <assert.h>
25 
26 #include <memory>
27 
28 #include "concurrency.h"
29 #include "aligned_allocator.h"
30 #include "container.h"
31 
32 static const int SLAB_LOCAL_BUF_SIZE = 100;
33 
34 class slab_allocator
35 {
36 public:
37  class linked_obj {
38  linked_obj *next;
39  public:
40  linked_obj() {
41  next = NULL;
42  }
43 
44  void add(linked_obj *obj) {
45  obj->next = this->next;
46  this->next = obj;
47  }
48 
49  linked_obj *get_next() {
50  return next;
51  }
52 
53  void set_next(linked_obj *obj) {
54  next = obj;
55  }
56  };
57 
58  class linked_obj_list {
59  linked_obj head;
60  linked_obj *end; // point to the last element or NULL.
61  int size;
62  public:
63  linked_obj_list() {
64  size = 0;
65  end = NULL;
66  }
67 
68  void add(linked_obj *obj) {
69  if (get_size() == 0) {
70  end = obj;
71  assert(head.get_next() == NULL);
72  }
73  head.add(obj);
74  size++;
75  }
76 
77  void add_list(linked_obj_list *list) {
78  if (list->get_size() == 0)
79  return;
80  if (get_size() == 0) {
81  size = list->get_size();
82  head = list->head;
83  end = list->end;
84  list->head.set_next(NULL);
85  list->end = NULL;
86  list->size = 0;
87  }
88  else {
89  list->end->set_next(head.get_next());
90  head = list->head;
91  size += list->size;
92  list->head.set_next(NULL);
93  list->end = NULL;
94  list->size = 0;
95  }
96  }
97 
98  /*
99  * It removes the specified number of objects from the list
100  * and returns them in a form of linked list..
101  */
102  linked_obj *pop(int num) {
103  linked_obj *obj = head.get_next();
104  int num_pop = 0;
105  linked_obj *prev_obj = NULL;
106  for (int i = 0; i < num && obj != NULL; i++) {
107  prev_obj = obj;
108  obj = obj->get_next();
109  num_pop++;
110  }
111  linked_obj *ret = head.get_next();
112  if (prev_obj)
113  prev_obj->set_next(NULL);
114 
115  head.set_next(obj);
116  size -= num_pop;
117  assert(size >= 0);
118  if (size == 0)
119  end = NULL;
120  return ret;
121  }
122 
123  int get_size() {
124  return size;
125  }
126  };
127 
128 private:
129  const int obj_size;
130  // the size to increase each time there aren't enough objects
131  const long increase_size;
132  // the maximal size of all objects
133  const long max_size;
134  const int node_id;
135  const int local_buf_size;
136  const bool thread_safe;
137 
138  linked_obj_list list;
139  // the current size of memory used by the allocator.
140  atomic_number<long> curr_size;
141  bool init;
142  bool pinned;
143 
144  std::vector<char *> alloc_bufs;
145 
146  pthread_spinlock_t lock;
147  // The buffers pre-allocated to serve allocation requests
148  // from the local threads.
149  pthread_key_t local_buf_key;
150 
151  thread_safe_FIFO_queue<fifo_queue<char *> *> per_thread_queues;
152 
153  std::string name;
154  static atomic_integer alloc_counter;
155 
156  fifo_queue<char *> *get_local_buf() {
157  fifo_queue<char *> *local_buf_refs
158  = (fifo_queue<char *> *) pthread_getspecific(local_buf_key);
159  if (local_buf_refs == NULL) {
160  local_buf_refs = fifo_queue<char *>::create(node_id, local_buf_size);
161  per_thread_queues.add(&local_buf_refs, 1);
162  pthread_setspecific(local_buf_key, local_buf_refs);
163  }
164  return local_buf_refs;
165  }
166 #ifdef MEMCHECK
167  aligned_allocator allocator;
168 #endif
169 public:
170  slab_allocator(const std::string &name, int _obj_size, long _increase_size,
171  // We allow pages to be pinned when allocated.
172  long _max_size, int _node_id, bool init = false, bool pinned = false,
173  int _local_buf_size = SLAB_LOCAL_BUF_SIZE, bool thread_safe = true);
174 
175  virtual ~slab_allocator();
176 
177  int get_obj_size() const {
178  return obj_size;
179  }
180 
181  int alloc(char **objs, int num);
182 
183  void free(char **objs, int num);
184 
185  char *alloc();
186 
187  void free(char *obj);
188 
189  // Use it carefully. It's not thread-safe.
190  bool contains(const char *addr) const {
191  for (unsigned i = 0; i < alloc_bufs.size(); i++) {
192  if (addr >= alloc_bufs[i] && addr < alloc_bufs[i]
193  + increase_size)
194  return true;
195  }
196  return false;
197  }
198 
199  long get_max_size() const {
200  return max_size;
201  }
202 
203  long get_curr_size() const {
204  return curr_size.get();
205  }
206 
207  const std::string &get_name() const {
208  return name;
209  }
210 };
211 
212 template<class T>
213 class obj_initiator
214 {
215 public:
216  typedef typename std::unique_ptr<obj_initiator<T> > ptr;
217  virtual void init(T *obj) = 0;
218 };
219 
220 template<class T>
221 class obj_destructor
222 {
223 public:
224  typedef typename std::unique_ptr<obj_destructor<T> > ptr;
225  virtual void destroy(T *obj) = 0;
226 };
227 
228 template<class T>
229 class default_obj_initiator: public obj_initiator<T>
230 {
231 public:
232  void init(T *obj) {
233  T tmp;
234  *obj = tmp;
235  }
236 };
237 
238 template<class T>
239 class default_obj_destructor: public obj_destructor<T>
240 {
241 public:
242  void destroy(T *obj) {
243  }
244 };
245 
246 template<class T>
247 class obj_allocator: public slab_allocator
248 {
249  typename obj_initiator<T>::ptr initiator;
250  typename obj_destructor<T>::ptr destructor;
251 public:
252  obj_allocator(const std::string &name, int node_id, bool thread_safe,
253  long increase_size, long max_size = INT_MAX,
254  typename obj_initiator<T>::ptr initiator = typename obj_initiator<T>::ptr(
255  new default_obj_initiator<T>()),
256  typename obj_destructor<T>::ptr destructor = typename obj_destructor<T>::ptr(
257  new default_obj_destructor<T>())
258  // leave some space for linked_obj, so the values in an object
259  // won't be modified.
260  ): slab_allocator(name, sizeof(T) + sizeof(slab_allocator::linked_obj),
261  increase_size, max_size, node_id, true, false, SLAB_LOCAL_BUF_SIZE,
262  thread_safe) {
263  assert(increase_size <= max_size);
264  this->initiator = std::move(initiator);
265  this->destructor = std::move(destructor);
266  }
267 
268  virtual int alloc_objs(T **objs, int num) {
269  char *addrs[num];
270  int ret = slab_allocator::alloc(addrs, num);
271  for (int i = 0; i < ret; i++) {
272  objs[i] = (T *) (addrs + sizeof(slab_allocator::linked_obj));
273  initiator->init(objs[i]);
274  }
275  return ret;
276  }
277 
278  virtual T *alloc_obj() {
279  char *addr = slab_allocator::alloc();
280  if (addr == NULL) {
281  return NULL;
282  }
283 
284  T *obj = (T *) (addr + sizeof(slab_allocator::linked_obj));
285  initiator->init(obj);
286  return obj;
287  }
288 
289  virtual void free(T **objs, int num) {
290  char *addrs[num];
291  for (int i = 0; i < num; i++) {
292  destructor->destroy(objs[i]);
293  addrs[i] = ((char *) objs[i]) - sizeof(slab_allocator::linked_obj);
294  }
295  slab_allocator::free(addrs, num);
296  }
297 
298  virtual void free(T *obj) {
299  char *addr = ((char *) obj) - sizeof(slab_allocator::linked_obj);
300  destructor->destroy(obj);
301  slab_allocator::free(addr);
302  }
303 };
304 
305 #endif