malloc.c

   1 /*
   2  * Allocates specified amount of memory.
   3  * Copyright (c) 2005,2007 by Michal Nazareicz <mina86@mina86.com>
   4  *
   5  * This program is free software; you can redistribute it and/or modify
   6  * it under the terms of the GNU General Public License as published by
   7  * the Free Software Foundation; either version 3 of the License, or
   8  * (at your option) any later version.
   9  *
  10  * This program is distributed in the hope that it will be useful,
  11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13  * GNU General Public License for more details.
  14  *
  15  * You should have received a copy of the GNU General Public License
  16  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  17  *
  18  * This is part of Tiny Applications Collection
  19  *   -> http://tinyapps.sourceforge.net/
  20  */
  21
  22 /*
  23  * This may be used to try to free some memory allocated for I/O
  24  * buffers (which could lead to some flushes as well), as well as it
  25  * can be used for testing the system performance when lots of memory
  26  * is consumed.
  27  */
  28
  29
  30 /********** Config **********/
  31
  32 /* Comment the next line out if your platform doesn't have sbrk()
  33    function */
  34 #define HAVE_SBRK
  35
  36
  37 /********** Includes **********/
  38 #define _BSD_SOURCE
  39 #define _SVID_SOURCE
  40 #define _DEFAULT_SOURCE
  41 #define _XOPEN_SOURCE 500
  42
  43 #include <stdlib.h>
  44 #include <stdio.h>
  45 #include <signal.h>
  46 #include <string.h>
  47 #ifdef HAVE_SBRK
  48 # include <errno.h>
  49 # include <unistd.h>
  50 #endif
  51
  52
  53 /********** Variables **********/
  54 static volatile sig_atomic_t signum = 0;
  55
  56
  57 /********** Function declarations **********/
  58 static unsigned long parse_arg(char *arg);
  59 static void usage(void);
  60 static void handle_all(void (*handler)(int signum));
  61 static void signal_handler(int sig);
  62 static int  alloc(unsigned long num);
  63 static void progress(unsigned long allocated, unsigned long to_allocate);
  64
  65 static unsigned long min(unsigned long x, unsigned long y) {
  66         return x < y ? x : y;
  67 }
  68
  69
  70 /********** Main **********/
  71 int main(int argc, char **argv) {
  72         unsigned long to_allocate, allocated = 0, dot = 1024;
  73
  74         switch (argc) {
  75         case 3:
  76                 if (!strcmp(argv[1], "-w")) {
  77                         ++argv;
  78                         /* FALL THROUGH */
  79         case 2:
  80                         to_allocate = parse_arg(argv[1]);
  81                         if (to_allocate) {
  82                                 break;
  83                         }
  84                 }
  85         default:
  86                 usage();
  87                 return 2;
  88         }
  89
  90         handle_all(signal_handler);
  91
  92         if (to_allocate < dot) {
  93                 dot = 4;
  94         }
  95
  96         progress(0, to_allocate);
  97         while (!signum &&
  98                allocated < to_allocate &&
  99                alloc(min(dot, to_allocate - allocated))) {
 100                 allocated += min(dot, to_allocate - allocated);
 101                 progress(allocated, to_allocate);
 102         }
 103         putchar('\n');
 104
 105         handle_all(SIG_DFL);
 106
 107         if (argc == 3) {
 108                 char buffer[1024];
 109                 while (fgets(buffer, sizeof buffer, stdin) && !strchr(buffer, '\n'))
 110                         /* nop */;
 111         }
 112
 113         return signum ? -signum : (allocated < to_allocate ? 1 : 0);
 114 }
 115
 116
 117 /********** Command line arguments **********/
 118 static unsigned long parse_arg(char *arg) {
 119         double ret = strtod(arg, &arg);
 120         switch (*arg) {
 121         case 'K':               ++arg; break;
 122         case 'M': ret *= 1<<10; ++arg; break;
 123         case 'G': ret *= 1<<20; ++arg; break;
 124         }
 125         return *arg ? 0 : ret;
 126 }
 127
 128 static void usage(void) {
 129         puts("usage: malloc [ -w ] <bytes>\n"
 130              " -w       wait for new line after allocating\n"
 131              " <bytes>  number of bytes to allocates;\n"
 132              "          can be fallowed by K (the default), M or G\n");
 133 }
 134
 135
 136 /********** Signal handler **********/
 137 static void signal_handler(int sig) {
 138         signum = sig;
 139         signal(sig, signal_handler);
 140 }
 141
 142 static void handle_all(void (*handler)(int signum)) {
 143         int i = 31;
 144         do {
 145                 switch (i) {
 146                 case SIGCONT: /* Continue if stopped */
 147                 case SIGSTOP: /* Stop process */
 148                 case SIGTSTP: /* Stop typed at tty */
 149                 case SIGTTIN: /* tty input for background process */
 150                 case SIGTTOU: /* tty output for background process */
 151                         break;
 152
 153                 default:
 154                         signal(i, handler);
 155                 }
 156         } while (--i);
 157 }
 158
 159
 160 /********** Allocates memory **********/
 161 static int alloc(unsigned long num) {
 162         char *ptr;
 163 #ifdef HAVE_SBRK
 164         errno = 0;
 165         ptr = sbrk(num <<= 10);
 166         if (errno) {
 167                 return 0;
 168         }
 169 #else
 170         ptr = malloc(num <<= 10);
 171         if (!ptr) {
 172                 return 0;
 173         }
 174 #endif
 175         do {
 176                 *ptr++ = --num;
 177         } while (!signum && num);
 178         return !signum;
 179 }
 180
 181
 182 /********** Prints progress **********/
 183 static void progress(unsigned long allocated, unsigned long to_allocate) {
 184         /* Zi an Yi are unofficial; Yi is 2^80 so there is no way we will
 185            ever allocat that amout of memoty ;) therefore we don't need to
 186            check whether there are some units available */
 187         static const char *const units = "KMGTPEZY";
 188
 189         static char dots[] = "=============================="
 190                 "==============================>";
 191
 192         static unsigned long old_size = ~0UL, old = ~0UL;
 193         static unsigned old_unit = ~0U;
 194
 195         unsigned long size = allocated;
 196         unsigned i;
 197
 198         /* Format size in a friendly way */
 199         for (i = 0; size > 102400; size >>= 10) {
 200                 ++i;
 201         }
 202         if (old_size == size && old_unit == i) {
 203                 return;
 204         }
 205
 206         old_size = size;
 207         old_unit = i;
 208
 209         if (size < 1024) {
 210                 printf("\r%5lu %ciB", size, units[i]);
 211         } else {
 212                 printf("\r%3lu.%lu %ciB", size / 1024, size * 10 / 1024 % 10,
 213                            units[i + 1]);
 214         }
 215
 216         /* Percentage */
 217         printf(" (%3lu%%)", allocated * 100 / to_allocate);
 218
 219         /* Slider */
 220         if ((allocated * 100 / to_allocate) != old * 100 / to_allocate) {
 221                 old = allocated;
 222
 223                 i = allocated * 50 / to_allocate;
 224                 dots[i] = '>';
 225                 printf("  [%-50.*s]", i + (allocated != to_allocate), dots);
 226                 dots[i] = '=';
 227         }
 228
 229         putchar('\r');
 230         fflush(stdout);
 231 }