code / lua-sampleprof

   1 /*
   2 * lsampleprof.c
   3 * A Lua sample profiling library by Bruce Hill. This library returns a single
   4 * function of the form profile([backprop=.619, [rate=100]], fn, ...)
   5 * that runs `fn` with any extra args, performs sample profiling, and
   6 * returns a table of the results.
   7 *
   8 *    - `rate` the average number of samples per millisecond
   9 *    - `backprop` controls how weights are propagated up the callstack. 0 means "only tally
  10 *       the current line", 1 means "tally the current line and everything above it in the
  11 *       callstack equally", 0.5 means "tally each line with half weight of the thing below
  12 *       it in the callstack"
  13 *  
  14 *  The results table is a mapping from "file:line" -> number of times a random sample landed on that line
  15 */
  16 
  17 #include "lua.h"
  18 #include "lauxlib.h"
  19 
  20 #include <unistd.h>
  21 #include <stdio.h>
  22 #include <stdlib.h>
  23 #include <string.h>
  24 #include <signal.h>
  25 #include <math.h>
  26 
  27 #define DEFAULT_BACKPROP 0.619
  28 #define DEFAULT_RATE 100
  29 
  30 static lua_State *lastL;
  31 static lua_Number backprop = DEFAULT_BACKPROP;
  32 static lua_Number rate = DEFAULT_RATE;
  33 
  34 static inline void randomalarm()
  35 {
  36     // Exponential decay:
  37     int r = rand();
  38     int delay = -log((double)r/(double)RAND_MAX) * 1e3 / rate;
  39     ualarm(1 + delay, 0);
  40 }
  41 
  42 static inline void bump(lua_State *L, lua_Debug *ar, lua_Number amount)
  43 {
  44     int profile_index = lua_gettop(L);
  45     lua_getinfo(L, "lS", ar);
  46     lua_pushfstring(L, "%s:%d", ar->short_src, ar->currentline);
  47     int key_index = lua_gettop(L);
  48     lua_pushvalue(L, key_index);
  49     int type = lua_gettable(L, profile_index);
  50     if (type == LUA_TNIL) {
  51         lua_pop(L, 1);
  52         lua_pushnumber(L, amount);
  53     } else {
  54         lua_Number count = lua_tonumber(L, -1);
  55         count += amount;
  56         lua_pop(L, 1);
  57         lua_pushnumber(L, count);
  58     }
  59     lua_settable(L, profile_index);
  60 }
  61 
  62 static void ltakesample(lua_State *L, lua_Debug *ar)
  63 {
  64     lua_pushlightuserdata(L, &lastL);
  65     lua_gettable(L, LUA_REGISTRYINDEX);
  66 
  67     lua_Number weight = 1.0;
  68     bump(L, ar, weight);
  69     lua_Debug ar2;
  70     for (int i = 0; lua_getstack(L, i, &ar2); i++) {
  71         weight *= backprop;
  72         bump(L, &ar2, weight);
  73     }
  74     // Disable hook and set alarm
  75     lua_sethook(lastL, NULL, 0, 0);
  76     randomalarm();
  77 }
  78 
  79 static void handler(int sig)
  80 {
  81     (void)sig;
  82     lua_sethook(lastL, ltakesample, LUA_MASKLINE, 0);
  83 }
  84 
  85 static int Lwrap(lua_State *L)
  86 {
  87     // Backprop
  88     if (lua_type(L, 1) == LUA_TNUMBER) {
  89         backprop = lua_tonumber(L, 1);
  90         lua_remove(L, 1);
  91     }
  92     // Sampling rate
  93     if (lua_type(L, 1) == LUA_TNUMBER) {
  94         rate = lua_tonumber(L, 1);
  95         lua_remove(L, 1);
  96     }
  97 
  98     struct sigaction newaction, oldaction;
  99     memset(&newaction, sizeof(newaction), 0);
 100     newaction.sa_handler = handler;
 101     newaction.sa_flags = SA_NODEFER;
 102     sigaction(SIGALRM, &newaction, &oldaction);
 103     randomalarm();
 104 
 105     lastL = L;
 106     lua_call(L, lua_gettop(L)-1, 0);
 107     lua_sethook(L, NULL, 0, 0);
 108 
 109     alarm(0);
 110     //sigaction(SIGALRM, &oldaction, NULL);
 111 
 112     backprop = DEFAULT_BACKPROP;
 113     rate = DEFAULT_RATE;
 114 
 115     lua_pushlightuserdata(L, &lastL);
 116     lua_gettable(L, LUA_REGISTRYINDEX);
 117 
 118     return 1;
 119 }
 120 
 121 LUALIB_API int luaopen_sampleprof(lua_State *L)
 122 {
 123     lua_pushlightuserdata(L, &lastL);
 124     lua_createtable(L, 0, 128);
 125     lua_settable(L, LUA_REGISTRYINDEX);
 126     lua_pushcfunction(L, Lwrap);
 127     return 1;
 128 }