0744f3f733ef031cc8d4502fa22a95ec2a5683cf
[rahunas] / src / rh-task-bandwidth.c
1 /**
2  * RahuNAS task bandwidth implementation 
3  * Author: Neutron Soutmun <neo.neutron@gmail.com>
4  * Date:   2008-11-20
5  */
6
7 #include <stdlib.h>
8 #include <syslog.h>
9 #include <sys/types.h>
10 #include <unistd.h>
11 #include <sys/wait.h>
12 #include <sys/stat.h>
13 #include <fcntl.h>
14
15 #include "rahunasd.h"
16 #include "rh-task.h"
17 #include "rh-task-bandwidth.h"
18 #include "rh-task-memset.h"
19 #include "rh-utils.h"
20
21
22
23 #define BANDWIDTH_WRAPPER "/etc/rahunas/bandwidth.sh"
24
25 static unsigned short slot_flags[MAX_SLOT_PAGE] = {1};
26 static unsigned short slot_count = 0;
27 static int bw_service = 0;
28
29 unsigned short _get_slot_id()
30 {
31   unsigned short slot_id    = 0;
32   unsigned short page       = 0;
33   unsigned char  id_on_page = 0;
34   time_t random_time;
35
36   // Slot is full 
37   if (slot_count >= MAX_SLOT_ID)
38     return 0;
39
40   // Do a random slot_id 
41   while (slot_id == 0) {
42     time(&(random_time));
43     srandom(random_time);
44     slot_id = random()/(int)(((unsigned int)RAND_MAX + 1) / (MAX_SLOT_ID + 1));
45    
46     // Check validity
47     page = slot_id / PAGE_SIZE; 
48     id_on_page = slot_id % PAGE_SIZE;
49    
50     if (!(slot_flags[page] & (1 << id_on_page))) {
51       // Slot available
52       break;
53     }
54
55     // Slot not available, retry
56     slot_id = 0;
57   }
58
59   return slot_id;
60 }
61
62 void mark_reserved_slot_id(unsigned int slot_id)
63 {
64   unsigned short page       = 0;
65   unsigned char  id_on_page = 0;
66
67   page = slot_id / PAGE_SIZE; 
68   id_on_page = slot_id % PAGE_SIZE;
69
70   slot_count++;
71   slot_flags[page] |= 1 << id_on_page;
72 }
73
74 int bandwidth_exec(struct vserver *vs, char *const args[])
75 {
76   pid_t ws;
77   pid_t pid;
78   int status;
79   int exec_pipe[2];
80   char buffer[150];
81   char *endline = NULL;
82   int ret = 0;
83   int fd = 0;
84   
85   memset(buffer, '\0', sizeof(buffer));
86
87   if (pipe(exec_pipe) == -1) {
88     logmsg(RH_LOG_ERROR, "Error: pipe()");
89     return -1;
90   }
91   DP("pipe0=%d,pipe1=%d", exec_pipe[0], exec_pipe[1]);
92
93   pid = vfork();
94   dup2(exec_pipe[1], STDOUT_FILENO);
95
96   if (pid == 0) {
97     // Child
98     execv(BANDWIDTH_WRAPPER, args);
99   } else if (pid < 0) {
100     // Fork error
101     logmsg(RH_LOG_ERROR, "Error: vfork()"); 
102     ret = -1;
103   } else {
104     // Parent
105     ws = waitpid(pid, &status, 0);
106
107     DP("Bandwidth: Return (%d)", WEXITSTATUS(status));
108
109     // Return message log
110     DP("Read message");
111     read(exec_pipe[0], buffer, sizeof(buffer));
112
113     if (buffer != NULL) {
114       DP("Got message: %s", buffer);
115       endline = strstr(buffer, "\n");
116       if (endline != NULL) 
117         *endline = '\0';
118
119       if (vs != NULL) {
120         logmsg(RH_LOG_NORMAL, "[%s] Bandwidth: %s", 
121           vs->vserver_config->vserver_name, buffer);
122       } else {
123         logmsg(RH_LOG_NORMAL, "[main server] Bandwidth: %s", buffer);
124       }
125     }
126
127     if (WIFEXITED(status)) {
128       ret = WEXITSTATUS(status);
129     } else {
130       ret = -1;
131     } 
132   }
133
134   close(exec_pipe[0]);
135   close(exec_pipe[1]);
136   return ret;
137 }
138
139 int bandwidth_start()
140 {
141   char *args[3];
142
143   DP("Bandwidth: start");
144
145   args[0] = BANDWIDTH_WRAPPER;
146   args[1] = "start";
147   args[2] = (char *) 0;
148
149   return bandwidth_exec(NULL, args);
150 }
151
152 int bandwidth_stop()
153 {
154   char *args[3];
155
156   DP("Bandwidth: stop");
157
158   args[0] = BANDWIDTH_WRAPPER;
159   args[1] = "stop";
160   args[2] = (char *) 0;
161
162   return bandwidth_exec(NULL, args);
163 }
164
165 int bandwidth_add(struct vserver *vs, struct bandwidth_req *bw_req)
166 {
167   char *args[7];
168
169   DP("Bandwidth: request %s %s %s %s", bw_req->slot_id, 
170      bw_req->ip, bw_req->bandwidth_max_down, bw_req->bandwidth_max_up);
171
172   args[0] = BANDWIDTH_WRAPPER;
173   args[1] = "add";
174   args[2] = bw_req->slot_id;
175   args[3] = bw_req->ip;
176   args[4] = bw_req->bandwidth_max_down;
177   args[5] = bw_req->bandwidth_max_up;
178   args[6] = (char *) 0;
179
180   return bandwidth_exec(vs, args);
181 }
182
183 int bandwidth_del(struct vserver *vs, struct bandwidth_req *bw_req)
184 {
185   char *args[4];
186
187   DP("Bandwidth: request %s", bw_req->slot_id);
188
189   args[0] = BANDWIDTH_WRAPPER;
190   args[1] = "del";
191   args[2] = bw_req->slot_id;
192   args[3] = (char *) 0;
193
194   return bandwidth_exec(vs, args);
195 }
196
197 /* Initialize */
198 static void init (struct vserver *vs)
199 {
200   logmsg(RH_LOG_NORMAL, "[%s] Task BANDWIDTH init..", 
201          vs->vserver_config->vserver_name);  
202 }
203
204 /* Cleanup */
205 static int cleanup (struct vserver *vs)
206 {
207   logmsg(RH_LOG_NORMAL, "[%s] Task BANDWIDTH cleanup..",
208          vs->vserver_config->vserver_name);  
209 }
210
211 /* Start service task */
212 static int startservice (struct vserver *vs)
213 {
214   if (!(bw_service++))
215     return bandwidth_start();
216
217   return 0;
218 }
219
220 /* Stop service task */
221 static int stopservice  (struct vserver *vs)
222 {
223   if ((--bw_service) == 0)
224     return bandwidth_stop();
225   
226   return 0;
227 }
228
229 /* Start session task */
230 static int startsess (struct vserver *vs, struct task_req *req)
231 {
232   struct bandwidth_req bw_req;
233   unsigned short slot_id;
234   unsigned char max_try = 3;
235   GList *member_node = NULL;
236   struct rahunas_member *member = NULL;
237
238   member_node = member_get_node_by_id(vs, req->id);
239   if (member_node == NULL)
240     return (-1);
241
242   member = (struct rahunas_member *) member_node->data;
243
244   if (member->bandwidth_max_down == 0 && member->bandwidth_max_up == 0)
245     return 0;
246
247   if (member->bandwidth_slot_id > 0)
248     return 0;
249   
250   // Formating the bandwidth request
251   sprintf(bw_req.ip, "%s", idtoip(vs->v_map, req->id));
252   sprintf(bw_req.bandwidth_max_down, "%lu", 
253     member->bandwidth_max_down);
254   sprintf(bw_req.bandwidth_max_up, "%lu", 
255     member->bandwidth_max_up);
256   
257   while (max_try > 0) { 
258     slot_id = _get_slot_id();
259     sprintf(bw_req.slot_id, "%d", slot_id);
260     if (bandwidth_add(vs, &bw_req) == 0)
261       break;
262     else
263       max_try--;
264   }
265
266
267   if (max_try == 0) {
268     logmsg(RH_LOG_ERROR, "[%s] Bandwidth: Maximum trying, failed!",
269            vs->vserver_config->vserver_name);
270     return -1;
271   }
272
273   mark_reserved_slot_id(slot_id);
274   member->bandwidth_slot_id = slot_id;
275  
276   return 0;
277 }
278
279 /* Stop session task */
280 static int stopsess  (struct vserver *vs, struct task_req *req)
281 {
282   struct bandwidth_req bw_req;
283   unsigned short slot_id = 0;
284   GList *member_node = NULL;
285   struct rahunas_member *member = NULL;
286   
287   member_node = member_get_node_by_id(vs, req->id);
288   if (member_node == NULL)
289     return (-1); 
290   
291   member = (struct rahunas_member *) member_node->data;
292   slot_id = member->bandwidth_slot_id;
293
294   if (slot_id < 1)
295     return 0;
296
297   sprintf(bw_req.slot_id, "%d", slot_id);
298
299   if (bandwidth_del(vs, &bw_req) == 0) {
300     member->bandwidth_slot_id = 0;
301
302     if (slot_count > 0)
303       slot_count--;
304   
305     return 0; 
306   } else {
307     return -1;
308   }
309 }
310
311 /* Commit start session task */
312 static int commitstartsess (struct vserver *vs, struct task_req *req)
313 {
314   /* Do nothing or need to implement */
315 }
316
317 /* Commit stop session task */
318 static int commitstopsess  (struct vserver *vs, struct task_req *req)
319 {
320   /* Do nothing or need to implement */
321 }
322
323 /* Rollback start session task */
324 static int rollbackstartsess (struct vserver *vs, struct task_req *req)
325 {
326   /* Do nothing or need to implement */
327 }
328
329 /* Rollback stop session task */
330 static int rollbackstopsess  (struct vserver *vs, struct task_req *req)
331 {
332   /* Do nothing or need to implement */
333 }
334
335 static struct task taskbandwidth = {
336   .taskname = "BANDWIDTH",
337   .taskprio = 20,
338   .init = &init,
339   .cleanup = &cleanup,
340   .startservice = &startservice,
341   .stopservice = &stopservice,
342   .startsess = &startsess,
343   .stopsess = &stopsess,
344   .commitstartsess = &commitstartsess,
345   .commitstopsess = &commitstopsess,
346   .rollbackstartsess = &rollbackstartsess,
347   .rollbackstopsess = &rollbackstopsess,
348 };
349
350 void rh_task_bandwidth_reg(struct main_server *ms) {
351   task_register(ms, &taskbandwidth);
352 }