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);

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.