Skip to content

Instantly share code, notes, and snippets.

@maxlapshin

maxlapshin/shmem.c

Created July 11, 2024 17:13
Show Gist options
  • Save maxlapshin/c0ee04bee5c075dab93c95575e66ead4 to your computer and use it in GitHub Desktop.
Save maxlapshin/c0ee04bee5c075dab93c95575e66ead4 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
shmem.c
#include <sys/mman.h>
#include <sys/stat.h>
#include <erl_nif.h>
#include <unistd.h>
#include <stdint.h>
#include <string.h>
#include <fcntl.h>
static ErlNifResourceType* shmem_resource;
struct Shmem {
uint8_t *base;
size_t len;
int can_write;
int shm;
char shm_name[1024];
};
static void
shmem_destructor(ErlNifEnv* env, void* obj) {
struct Shmem *s = (struct Shmem *)obj;
if(s->base) {
munmap(s->base, s->len);
close(s->shm);
}
s->base = 0;
s->shm = -1;
}
static int
load(ErlNifEnv* env, void** priv, ERL_NIF_TERM load_info) {
if (!shmem_resource) {
ErlNifResourceFlags flags = (ErlNifResourceFlags) (ERL_NIF_RT_CREATE | ERL_NIF_RT_TAKEOVER);
shmem_resource = enif_open_resource_type(env, NULL, (char *)"shmem_resource", &shmem_destructor, flags, NULL);
}
return 0;
}
static int
upgrade(ErlNifEnv* env, void** priv, void** old_priv, ERL_NIF_TERM load_info) {
return 0;
}
static int
reload(ErlNifEnv* env, void** priv, ERL_NIF_TERM load_info) {
return 0;
}
static void
unload(ErlNifEnv* env, void* priv) {
return;
}
static ERL_NIF_TERM
shmem_open(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
ErlNifUInt64 len;
ErlNifBinary shm_bin;
if(!enif_inspect_binary(env, argv[0], &shm_bin)) {
return enif_raise_exception(env, enif_make_atom(env, "bad_shm_name"));
}
if(!enif_get_uint64(env, argv[1], &len)) {
return enif_raise_exception(env, enif_make_atom(env, "bad_len"));
}
int write_flag = enif_make_atom(env, "true") == argv[2];
int create_flag = enif_make_atom(env, "true") == argv[3];
int flags = write_flag ? O_RDWR : O_RDONLY;
if(create_flag) {
flags |= O_CREAT;
}
char shm_name[1024];
memset(shm_name, 0, sizeof(shm_name));
strncpy(shm_name, (const char *)shm_bin.data, sizeof(shm_name) > shm_bin.size ? shm_bin.size : sizeof(shm_name));
int shm = shm_open(shm_name, flags, S_IRUSR | S_IWUSR);
if(shm < 0) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "shm_open_failed")
);
}
struct stat statbuf;
fstat(shm, &statbuf);
if(statbuf.st_size < len) {
ftruncate(shm, len);
}
uint8_t *buffer = mmap(NULL, len, PROT_READ | (write_flag ? PROT_WRITE : 0), MAP_SHARED, shm, 0);
if(!buffer) {
close(shm);
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "mmap_failed")
);
}
struct Shmem *shmem = (struct Shmem *)enif_alloc_resource(shmem_resource, sizeof(struct Shmem));
shmem->base = buffer;
shmem->len = len;
shmem->shm = shm;
shmem->can_write = write_flag;
memcpy(shmem->shm_name, shm_name, sizeof(shm_name));
ERL_NIF_TERM shmem_term = enif_make_resource_binary(env, shmem, shmem->shm_name, strlen(shmem->shm_name));
enif_release_resource(shmem);
return enif_make_tuple2(env, enif_make_atom(env, "ok"), shmem_term);
}
static ERL_NIF_TERM
shmem_close(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
struct Shmem *shmem;
if (!enif_get_resource(env, argv[0], shmem_resource, (void **)&shmem)) {
return enif_raise_exception(env, enif_make_atom(env, "not_a_shmem"));
}
if(shmem->base) {
munmap(shmem->base, shmem->len);
close(shmem->shm);
shmem->base = 0;
shmem->shm = -1;
}
return enif_make_atom(env, "ok");
}
static ERL_NIF_TERM
shmem_read(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
struct Shmem *shmem;
if (!enif_get_resource(env, argv[0], shmem_resource, (void **)&shmem)) {
return enif_raise_exception(env, enif_make_atom(env, "not_a_shmem"));
}
if(!shmem->base) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "closed")
);
}
ErlNifSInt64 offset;
ErlNifSInt64 len;
if (!enif_get_int64(env, argv[1], &offset)) {
return enif_raise_exception(env, enif_make_atom(env, "not_an_offset"));
}
if (!enif_get_int64(env, argv[2], &len)) {
return enif_raise_exception(env, enif_make_atom(env, "not_a_len"));
}
if(offset < 0 || len <= 0 || offset + len > shmem->len) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "out_of_range")
);
}
ErlNifBinary bin;
if(!enif_alloc_binary(len, &bin)) {
return enif_raise_exception(env, enif_make_atom(env, "alloc"));
}
memcpy(bin.data, shmem->base + offset, len);
return enif_make_tuple2(env,
enif_make_atom(env, "ok"),
enif_make_binary(env, &bin)
);
}
static ERL_NIF_TERM
shmem_write(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
struct Shmem *shmem;
if (!enif_get_resource(env, argv[0], shmem_resource, (void **)&shmem)) {
return enif_raise_exception(env, enif_make_atom(env, "not_a_shmem"));
}
if(!shmem->base) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "closed")
);
}
if(!shmem->can_write) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "readonly")
);
}
ErlNifSInt64 offset;
ErlNifBinary bin;
if (!enif_get_int64(env, argv[1], &offset)) {
return enif_raise_exception(env, enif_make_atom(env, "not_an_offset"));
}
if(!enif_inspect_binary(env, argv[2], &bin)) {
return enif_raise_exception(env, enif_make_atom(env, "not_a_bin"));
}
if(offset < 0 || offset + bin.size > shmem->len) {
return enif_make_tuple2(env,
enif_make_atom(env, "error"),
enif_make_atom(env, "out_of_range")
);
}
memcpy(shmem->base + offset, bin.data, bin.size);
return enif_make_atom(env, "ok");
}
static ErlNifFunc funcs[] = {
{"open0", 4, shmem_open, 0},
{"close", 1, shmem_close, 0},
{"pread", 3, shmem_read, 0},
{"pwrite", 3, shmem_write, 0}
};
ERL_NIF_INIT(shmem, funcs, load, reload, upgrade, unload);
Raw
shmem.erl
-module(shmem).
-on_load(load_nif/0).
-export([open/2, close/1, pread/3, pwrite/3]).
load_nif() ->
Path = case code:lib_dir(corelib) of
P when is_list(P) -> filename:join(P, priv);
_ -> "./priv"
end,
% check if there is a file on disk
NifName = Path ++ "/shmem",
case file:read_file_info(NifName ++ ".so") of
{ok, _} ->
case erlang:load_nif(NifName, LoadInfo) of
ok ->
ok;
{error, {load_failed, NifLoadError}} ->
case re:run(NifLoadError, "'([^']+)'",[{capture,all_but_first,list}]) of
{match, [RealNifMsg]} -> {error, {load_failed, RealNifMsg}};
_ -> {error, {load_failed, NifLoadError}}
end;
{error, NifLoadError} ->
{error, NifLoadError}
end;
{error, enoent} -> ok
end.
-type shmem() :: any().
-type shmem_opts() :: #{
len := pos_integer(),
write => true,
create => true
}.
-type shm_error() ::
out_of_range |
readonly |
mmap_failed |
shm_open_failed |
closed.
-spec open(ShmName, Opts) -> {ok, shmem()} | {error, shm_error()}
when ShmName :: binary(), Opts :: shmem_opts().
open(ShmName, #{len := Length} = Opts) ->
Write = maps:get(write, Opts, false),
Create = maps:get(create, Opts, false),
open0(ShmName, Length, Write, Create).
open0(ShmName, Len, Write, Create) ->
erlang:nif_error(not_loaded, [ShmName, Len, Write, Create]).
-spec pread(Shm, Offset, Len) -> {ok, binary()} | {error, shm_error()}
when Shm::shmem(), Offset::non_neg_integer(), Len::pos_integer().
pread(Shm, Offset, Len) ->
erlang:nif_error(not_loaded, [Shm, Offset, Len]).
-spec pwrite(Shm, Offset, Binary) -> ok | {error, shm_error()}
when Shm::shmem(), Offset::non_neg_integer(), Binary::binary().
pwrite(Shm, Offset, Binary) ->
erlang:nif_error(not_loaded, [Shm, Offset, size(Binary)]).
-spec close(Shm) -> ok
when Shm::shmem().
close(_Shm) ->
ok.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment