├── .gitignore ├── slowjit.control ├── Makefile ├── slowjit.c ├── README.md ├── slowjit.h └── slowjit_expr.c /.gitignore: -------------------------------------------------------------------------------- 1 | build/* 2 | *.o 3 | *.so 4 | compile_commands.json 5 | .cache/* 6 | -------------------------------------------------------------------------------- /slowjit.control: -------------------------------------------------------------------------------- 1 | comment = 'A very inefficient jit provider.' 2 | default_version = '1.0.0' 3 | module_pathname = '$libdir/slowjit' 4 | relocatable = true 5 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | MODULE_big = slowjit 2 | EXTENSION = slowjit 3 | 4 | OBJS = slowjit.o slowjit_expr.o 5 | 6 | # Disable the bitcode generation. 7 | override with_llvm = no 8 | 9 | PG_CONFIG := pg_config 10 | PGXS := $(shell $(PG_CONFIG) --pgxs) 11 | include $(PGXS) 12 | -------------------------------------------------------------------------------- /slowjit.c: -------------------------------------------------------------------------------- 1 | #include "slowjit.h" 2 | 3 | PG_MODULE_MAGIC; 4 | 5 | char *slowjit_cc_path = NULL; 6 | 7 | void _PG_init(void) { 8 | DefineCustomStringVariable( 9 | "slowjit.cc_path", 10 | "Sets the C compiler to be used on the remote server.", NULL, &slowjit_cc_path, 11 | "cc", PGC_SU_BACKEND, GUC_EXPLAIN, NULL, NULL, NULL); 12 | } 13 | 14 | void _PG_jit_provider_init(JitProviderCallbacks *cb) { 15 | cb->reset_after_error = slowjit_reset_after_error; 16 | cb->release_context = slowjit_release_context; 17 | cb->compile_expr = slowjit_compile_expr; 18 | } 19 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pg_slowjit 2 | 3 | > A simple demo to illustrate how to implement a JIT provider for PostgreSQL. `pg_slowjit` emits C codes and compile them to shared libraries in runtime. This project is inspired by [09 - Query Compilation & JIT Code Generation (CMU Advanced Databases / Spring 2023)](https://www.youtube.com/watch?v=eurwtUhY5fk). 4 | 5 | ## What does it support? 6 | 7 | Currently, it only supports jitting a few operators. You can use it to jit the query `SELECT 1;`. 8 | 9 | ``` 10 | postgres=# EXPLAIN (SETTINGS ON) SELECT 1; 11 | QUERY PLAN 12 | --------------------------------------------------------------------------------- 13 | Result (cost=0.00..0.01 rows=1 width=4) 14 | Settings: jit_above_cost = '0' 15 | JIT: 16 | Functions: 1 17 | Options: Inlining false, Optimization false, Expressions true, Deforming true 18 | (5 rows) 19 | ``` 20 | -------------------------------------------------------------------------------- /slowjit.h: -------------------------------------------------------------------------------- 1 | #ifndef _SLOWJIT_H_ 2 | #define _SLOWJIT_H_ 3 | 4 | #include "postgres.h" 5 | 6 | #include "executor/execExpr.h" 7 | #include "executor/tuptable.h" 8 | #include "fmgr.h" 9 | #include "jit/jit.h" 10 | #include "lib/stringinfo.h" 11 | #include "miscadmin.h" 12 | #include "nodes/execnodes.h" 13 | #include "nodes/pg_list.h" 14 | #include "port.h" 15 | #include "portability/instr_time.h" 16 | #include "storage/ipc.h" 17 | #include "utils/expandeddatum.h" 18 | #include "utils/guc.h" 19 | #include "utils/memutils.h" 20 | #include "utils/palloc.h" 21 | #include "utils/resowner.h" 22 | #include "utils/resowner_private.h" 23 | 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | extern char *slowjit_cc_path; 30 | 31 | extern void _PG_init(void); 32 | extern void _PG_jit_provider_init(JitProviderCallbacks *cb); 33 | 34 | typedef struct SlowJitHandle { 35 | /* Shared library path, for error reporting. */ 36 | char *shared_library_path; 37 | /* Handle of the shared library. */ 38 | void *handle; 39 | } SlowJitHandle; 40 | 41 | typedef struct SlowJitContext { 42 | JitContext base; 43 | StringInfoData code_holder; 44 | /* Is there any pending code that needs to be emitted */ 45 | bool compiled; 46 | /* # of objects emitted, used to generate non-conflicting names */ 47 | int counter; 48 | 49 | size_t module_generation; 50 | 51 | /* Handles of the compiled shared libraries. */ 52 | List *handles; 53 | } SlowJitContext; 54 | 55 | typedef struct SlowJitCompiledExprState { 56 | SlowJitContext *jit_ctx; 57 | const char *funcname; 58 | } SlowJitCompiledExprState; 59 | 60 | extern bool slowjit_compile_expr(ExprState *state); 61 | extern void slowjit_release_context(JitContext *jit_ctx); 62 | extern void slowjit_reset_after_error(void); 63 | 64 | #endif 65 | -------------------------------------------------------------------------------- /slowjit_expr.c: -------------------------------------------------------------------------------- 1 | #include "slowjit.h" 2 | 3 | static size_t slowjit_module_generation = 0; 4 | 5 | static SlowJitContext *slowjit_create_context(int jit_flags) { 6 | SlowJitContext *jit_ctx = NULL; 7 | MemoryContext oldcontext; 8 | 9 | ResourceOwnerEnlargeJIT(CurrentResourceOwner); 10 | 11 | jit_ctx = MemoryContextAllocZero(TopMemoryContext, sizeof(SlowJitContext)); 12 | jit_ctx->base.flags = jit_flags; 13 | 14 | /* We emit the whole C program to the TopMemoryContext. */ 15 | oldcontext = MemoryContextSwitchTo(TopMemoryContext); 16 | initStringInfo(&jit_ctx->code_holder); 17 | MemoryContextSwitchTo(oldcontext); 18 | 19 | /* ensure cleanup */ 20 | jit_ctx->base.resowner = CurrentResourceOwner; 21 | ResourceOwnerRememberJIT(CurrentResourceOwner, PointerGetDatum(jit_ctx)); 22 | 23 | return jit_ctx; 24 | } 25 | 26 | /* This function is copied and pasted from llvmjit. */ 27 | static char *slowjit_expand_funcname(SlowJitContext *jit_ctx, 28 | const char *basename) { 29 | jit_ctx->base.instr.created_functions++; 30 | return psprintf("%s_%zu_%d", basename, jit_ctx->module_generation, 31 | jit_ctx->counter++); 32 | } 33 | 34 | static void slowjit_compile_module(SlowJitContext *jit_ctx) { 35 | char shared_library_path[MAXPGPATH]; 36 | snprintf(shared_library_path, MAXPGPATH, "/tmp/%d.%zu.so", MyProcPid, 37 | jit_ctx->module_generation); 38 | { 39 | char c_src_path[MAXPGPATH]; 40 | /* Compile the program to a shared library. */ 41 | char include_server_path[MAXPGPATH]; 42 | FILE *c_src; 43 | char command[MAXPGPATH]; 44 | 45 | /* Write the C program to a dot C file. */ 46 | snprintf(c_src_path, MAXPGPATH, "/tmp/%d.%zu.c", MyProcPid, 47 | jit_ctx->module_generation); 48 | c_src = fopen(c_src_path, "w+"); 49 | if (c_src == NULL) { 50 | ereport(ERROR, (errmsg("cannot open file '%s' for write", c_src_path))); 51 | } 52 | fwrite(jit_ctx->code_holder.data, 1, jit_ctx->code_holder.len, c_src); 53 | fclose(c_src); 54 | 55 | /* Code holder now is useless. */ 56 | resetStringInfo(&jit_ctx->code_holder); 57 | 58 | /* Prepare compile command. */ 59 | get_includeserver_path(my_exec_path, include_server_path); 60 | snprintf(command, MAXPGPATH, "%s -fPIC -I%s -shared -ggdb -g3 -O0 -o %s %s", 61 | slowjit_cc_path, include_server_path, shared_library_path, 62 | c_src_path); 63 | if (system(command) != 0) { 64 | ereport(ERROR, (errmsg("cannot execute command: %s", command))); 65 | } 66 | } 67 | 68 | { 69 | /* Load the compiled shared library to the backend process. */ 70 | void *shared_library_handle; 71 | SlowJitHandle *slowjit_handle; 72 | MemoryContext oldcontext; 73 | 74 | shared_library_handle = dlopen(shared_library_path, RTLD_LAZY); 75 | if (shared_library_handle == NULL) { 76 | char *err = dlerror(); 77 | ereport(ERROR, 78 | (errmsg("cannot dlopen '%s': %s", shared_library_path, err))); 79 | } 80 | 81 | oldcontext = MemoryContextSwitchTo(TopMemoryContext); 82 | slowjit_handle = (SlowJitHandle *)palloc0(sizeof(SlowJitHandle)); 83 | slowjit_handle->handle = shared_library_handle; 84 | slowjit_handle->shared_library_path = pstrdup(shared_library_path); 85 | jit_ctx->handles = lappend(jit_ctx->handles, slowjit_handle); 86 | MemoryContextSwitchTo(oldcontext); 87 | } 88 | 89 | /* The current module is compiled. */ 90 | jit_ctx->compiled = true; 91 | } 92 | 93 | static ExprStateEvalFunc slowjit_get_function(SlowJitContext *jit_ctx, 94 | const char *funcname) { 95 | ListCell *lc; 96 | ExprStateEvalFunc jitted_func = NULL; 97 | 98 | if (!jit_ctx->compiled) { 99 | slowjit_compile_module(jit_ctx); 100 | } 101 | 102 | foreach (lc, jit_ctx->handles) { 103 | SlowJitHandle *slowjit_handle = (SlowJitHandle *)lfirst(lc); 104 | void *handle = slowjit_handle->handle; 105 | 106 | jitted_func = dlsym(handle, funcname); 107 | if (jitted_func == NULL) { 108 | /* Consume the existing error. */ 109 | char *err = dlerror(); 110 | ereport(LOG, (errmsg("cannot find symbol '%s' from '%s': %s", funcname, 111 | slowjit_handle->shared_library_path, err))); 112 | continue; 113 | } else { 114 | return jitted_func; 115 | } 116 | } 117 | 118 | ereport(ERROR, (errmsg("cannot jit function '%s'", funcname))); 119 | 120 | return NULL; 121 | } 122 | 123 | static Datum slowjit_exec_compiled_expr(ExprState *state, ExprContext *econtext, 124 | bool *isNull) { 125 | SlowJitCompiledExprState *cstate = state->evalfunc_private; 126 | ExprStateEvalFunc func; 127 | 128 | CheckExprStillValid(state, econtext); 129 | 130 | func = slowjit_get_function(cstate->jit_ctx, cstate->funcname); 131 | Assert(func); 132 | 133 | /* remove indirection via this function for future calls */ 134 | state->evalfunc = func; 135 | 136 | return func(state, econtext, isNull); 137 | } 138 | 139 | bool slowjit_compile_expr(ExprState *state) { 140 | PlanState *parent = state->parent; 141 | SlowJitContext *jit_ctx = NULL; 142 | char *funcname = NULL; 143 | StringInfoData jitted_expr_body; 144 | 145 | /* parent shouldn't be NULL. */ 146 | Assert(parent != NULL); 147 | 148 | /* Initialize the context. */ 149 | if (parent->state->es_jit) { 150 | jit_ctx = (SlowJitContext *)parent->state->es_jit; 151 | } else { 152 | jit_ctx = slowjit_create_context(parent->state->es_jit_flags); 153 | parent->state->es_jit = &jit_ctx->base; 154 | } 155 | 156 | initStringInfo(&jitted_expr_body); 157 | 158 | #define emit_line(...) \ 159 | do { \ 160 | appendStringInfo(&jitted_expr_body, __VA_ARGS__); \ 161 | appendStringInfoChar(&jitted_expr_body, '\n'); \ 162 | } while (0) 163 | 164 | #define emit_include(header) emit_line("#include \"%s\"", header) 165 | 166 | /* The code holder is empty, which means this is a new module. */ 167 | if (jit_ctx->code_holder.len == 0) { 168 | /* We only emit the header once! */ 169 | emit_include("postgres.h"); 170 | emit_include("nodes/execnodes.h"); 171 | 172 | jit_ctx->compiled = false; 173 | jit_ctx->module_generation = slowjit_module_generation++; 174 | } 175 | 176 | /* Emit the jitted function signature. */ 177 | funcname = slowjit_expand_funcname(jit_ctx, "slowjit_eval_expr"); 178 | emit_line("Datum %s(ExprState *state, ExprContext *econtext, bool " 179 | "*isnull)", 180 | funcname); 181 | 182 | /* Open function body. */ 183 | emit_line("{"); 184 | 185 | /* Emit some commonly used variables. */ 186 | emit_line(" TupleTableSlot *resultslot = state->resultslot;"); 187 | 188 | for (int opno = 0; opno < state->steps_len; ++opno) { 189 | ExprEvalStep *op; 190 | ExprEvalOp opcode; 191 | 192 | op = &state->steps[opno]; 193 | opcode = ExecEvalStepOp(state, op); 194 | 195 | switch (opcode) { 196 | case EEOP_DONE: { 197 | emit_line(" { // EEOP_DONE"); 198 | emit_line(" *isnull = state->resnull;"); 199 | emit_line(" }"); 200 | emit_line(" return state->resvalue;"); 201 | 202 | /* Close function boday. */ 203 | emit_line("}"); 204 | break; 205 | } 206 | case EEOP_ASSIGN_TMP: { 207 | emit_line(" { // EEOP_ASSIGN_TMP"); 208 | emit_line(" int resultnum = %d;", op->d.assign_tmp.resultnum); 209 | emit_line(" resultslot->tts_values[resultnum] = state->resvalue;"); 210 | emit_line(" resultslot->tts_isnull[resultnum] = state->resnull;"); 211 | emit_line(" }"); 212 | break; 213 | } 214 | case EEOP_CONST: { 215 | emit_line(" { // EEOP_CONST"); 216 | emit_line(" bool *resnull = (bool *) %lu;", (uint64_t)op->resnull); 217 | emit_line(" Datum *resvalue = (Datum *) %lu;", (uint64_t)op->resvalue); 218 | emit_line(" *resnull = (bool) %d;", op->d.constval.isnull); 219 | emit_line(" *resvalue = (Datum) %lu;", op->d.constval.value); 220 | emit_line(" }"); 221 | break; 222 | } 223 | default: { 224 | resetStringInfo(&jitted_expr_body); 225 | return false; 226 | } 227 | } 228 | } 229 | 230 | appendStringInfo(&jit_ctx->code_holder, "%s", jitted_expr_body.data); 231 | resetStringInfo(&jitted_expr_body); 232 | 233 | { 234 | SlowJitCompiledExprState *cstate = 235 | palloc0(sizeof(SlowJitCompiledExprState)); 236 | 237 | cstate->jit_ctx = jit_ctx; 238 | cstate->funcname = funcname; 239 | 240 | state->evalfunc = slowjit_exec_compiled_expr; 241 | state->evalfunc_private = cstate; 242 | } 243 | 244 | return true; 245 | } 246 | 247 | void slowjit_release_context(JitContext *jit_ctx) { 248 | SlowJitContext *ctx = (SlowJitContext *)jit_ctx; 249 | ListCell *lc; 250 | 251 | foreach (lc, ctx->handles) { 252 | SlowJitHandle *slowjit_handle = (SlowJitHandle *)lfirst(lc); 253 | void *handle = slowjit_handle->handle; 254 | dlclose(handle); 255 | pfree(slowjit_handle->shared_library_path); 256 | pfree(slowjit_handle); 257 | } 258 | 259 | list_free(ctx->handles); 260 | ctx->handles = NIL; 261 | } 262 | 263 | void slowjit_reset_after_error(void) { /* TODO */ 264 | } 265 | --------------------------------------------------------------------------------