maxlapshin · July 11, 2024 17:13 · Jul 11, 2024
diff --git a/shmem.c b/shmem.c
@@ -0,0 +1,239 @@
+#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);
+
diff --git a/shmem.erl b/shmem.erl
@@ -0,0 +1,80 @@
+-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.