erts: Add nif ioq

2 files changed, 421 insertions, 4 deletions

diff --git a/erts/emulator/test/nif_SUITE.erl b/erts/emulator/test/nif_SUITE.erl
index 0337274178..4811244b98 100644
--- a/erts/emulator/test/nif_SUITE.erl
+++ b/erts/emulator/test/nif_SUITE.erl
@@ -61,7 +61,8 @@
          nif_internal_hash_salted/1,
          nif_phash2/1,
          nif_whereis/1, nif_whereis_parallel/1,
-         nif_whereis_threaded/1, nif_whereis_proxy/1
+         nif_whereis_threaded/1, nif_whereis_proxy/1,
+         nif_ioq/1
 	]).
 
 -export([many_args_100/100]).
@@ -99,7 +100,8 @@ all() ->
      nif_internal_hash,
      nif_internal_hash_salted,
      nif_phash2,
-     nif_whereis, nif_whereis_parallel, nif_whereis_threaded].
+     nif_whereis, nif_whereis_parallel, nif_whereis_threaded,
+     nif_ioq].
 
 groups() ->
     [{G, [], api_repeaters()} || G <- api_groups()]
@@ -2957,6 +2959,180 @@ nif_whereis_proxy(Ref) ->
         {Ref, quit} ->
             ok
     end.
+nif_ioq(Config) ->
+    ensure_lib_loaded(Config),
+
+    Script =
+        [{create, a},
+
+         %% Test enq of erlang term binary
+         {enqb,   a},
+         {enqb,   a, 3},
+
+         %% Test enq of non-erlang term binary
+         {enqbraw,a},
+         {enqbraw,a, 5},
+         {peek,   a},
+         {deq,    a, 42},
+
+         %% Test enqv
+         {enqv,   a, 2, 100},
+         {deq,    a, all},
+
+         %% This skips all elements but one in the iolist
+         {enqv,   a, 5, iolist_size(nif_ioq_payload(5)) - 1},
+         {peek,   a},
+
+         %% Test to enqueue a bunch of refc binaries
+         {enqv,   a, [nif_ioq_payload(refcbin) || _ <- lists:seq(1,20)], 0},
+
+         %% Enq stuff to destroy with data in queue
+         {enqv,   a, 2, 100},
+         {destroy,a},
+
+         %% Test destroy of new queue
+         {create, a},
+         {destroy,a}
+        ],
+
+    nif_ioq_run(Script),
+
+    %% Test that only enif_inspect_as_vec works
+    Payload = nif_ioq_payload(5),
+    PayloadBin = iolist_to_binary(Payload),
+
+    [begin
+         PayloadBin = iolist_to_binary(ioq_nif(inspect,Payload,Stack,Env)),
+         <<>>       = iolist_to_binary(ioq_nif(inspect,[],Stack,Env))
+     end || Stack <- [no_stack, use_stack], Env <- [use_env, no_env]],
+
+    %% Test error cases
+
+    Q = ioq_nif(create),
+
+    {'EXIT', {badarg, _}} = (catch ioq_nif(deq, Q, 1)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, 1, 1234)),
+
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [atom_in_list], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [make_ref()], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [256], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [-1], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [#{}], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [1 bsl 64], 0)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(enqv, Q, [{tuple}], 0)),
+
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [atom_in_list], use_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [make_ref()], no_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [256], use_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [-1], no_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [#{}], use_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [1 bsl 64], no_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  [{tuple}], use_stack)),
+    {'EXIT', {badarg, _}} = (catch ioq_nif(inspect,  <<"binary">>, use_stack)),
+
+    ioq_nif(destroy, Q),
+
+    %% Test that the example in the docs works
+    ExampleQ = ioq_nif(create),
+    true = ioq_nif(example, ExampleQ, nif_ioq_payload(5)),
+    ioq_nif(destroy, ExampleQ),
+
+    ok.
+
+
+nif_ioq_run(Script) ->
+    nif_ioq_run(Script, #{}).
+
+nif_ioq_run([{Action, Name}|T], State)
+  when Action =:= enqb; Action =:= enqbraw ->
+    nif_ioq_run([{Action, Name, heapbin}|T], State);
+nif_ioq_run([{Action, Name, Skip}|T], State)
+  when Action =:= enqb, is_integer(Skip);
+       Action =:= enqbraw, is_integer(Skip) ->
+    nif_ioq_run([{Action, Name, heapbin, Skip}|T], State);
+nif_ioq_run([{Action, Name, N}|T], State)
+  when Action =:= enqv; Action =:= enqb; Action =:= enqbraw ->
+    nif_ioq_run([{Action, Name, N, 0}|T], State);
+nif_ioq_run([{Action, Name, N, Skip}|T], State)
+  when Action =:= enqv; Action =:= enqb; Action =:= enqbraw ->
+
+    #{ q := IOQ, b := B } = Q = maps:get(Name, State),
+    true = ioq_nif(size, IOQ) == iolist_size(B),
+
+    %% Sanitize the log output a bit so that it doesn't become too large.
+    H = {Action, Name, try iolist_size(N) of Sz -> Sz catch _:_ -> N end, Skip},
+    ct:log("~p", [H]),
+
+    Data = nif_ioq_payload(N),
+    ioq_nif(Action, IOQ, Data, Skip),
+
+    <<_:Skip/binary, SkippedData/binary>> = iolist_to_binary(Data),
+
+    true = ioq_nif(size, IOQ) == (iolist_size([B|SkippedData])),
+
+    nif_ioq_run(T, State#{ Name := Q#{ b := [B|SkippedData]}});
+nif_ioq_run([{peek, Name} = H|T], State) ->
+    #{ q := IOQ, b := B } = maps:get(Name, State),
+    true = ioq_nif(size, IOQ) == iolist_size(B),
+
+    ct:log("~p", [H]),
+
+    Data = ioq_nif(peek, IOQ, ioq_nif(size, IOQ)),
+
+    true = iolist_to_binary(B) == iolist_to_binary(Data),
+    nif_ioq_run(T, State);
+nif_ioq_run([{deq, Name, all}|T], State) ->
+    #{ q := IOQ, b := B } = maps:get(Name, State),
+    Size = ioq_nif(size, IOQ),
+    true = Size == iolist_size(B),
+    nif_ioq_run([{deq, Name, Size}|T], State);
+nif_ioq_run([{deq, Name, N} = H|T], State) ->
+    #{ q := IOQ, b := B } = Q = maps:get(Name, State),
+    true = ioq_nif(size, IOQ) == iolist_size(B),
+
+    ct:log("~p", [H]),
+
+    <<_:N/binary,Remain/binary>> = iolist_to_binary(B),
+    NewQ = Q#{ b := Remain },
+
+    Sz = ioq_nif(deq, IOQ, N),
+
+    true = Sz == iolist_size(Remain),
+    true = ioq_nif(size, IOQ) == iolist_size(Remain),
+
+    nif_ioq_run(T, State#{ Name := NewQ });
+nif_ioq_run([{create, Name} = H|T], State) ->
+    ct:log("~p", [H]),
+    nif_ioq_run(T, State#{ Name => #{ q => ioq_nif(create), b => [] } });
+nif_ioq_run([{destroy, Name} = H|T], State) ->
+    #{ q := IOQ, b := B } = maps:get(Name, State),
+    true = ioq_nif(size, IOQ) == iolist_size(B),
+
+    ct:log("~p", [H]),
+
+    ioq_nif(destroy, IOQ),
+
+    nif_ioq_run(T, maps:remove(Name, State));
+nif_ioq_run([], State) ->
+    State.
+
+nif_ioq_payload(N) when is_integer(N) ->
+    Tail = if N > 3 -> nif_ioq_payload(N-3); true -> [] end,
+    Head = element(1, lists:split(N,[nif_ioq_payload(subbin),
+                                     nif_ioq_payload(heapbin),
+                                     nif_ioq_payload(refcbin) | Tail])),
+    erlang:iolist_to_iovec(Head);
+nif_ioq_payload(subbin) ->
+    Bin = nif_ioq_payload(refcbin),
+    Sz = size(Bin) - 1,
+    <<_:8,SubBin:Sz/binary,_/bits>> = Bin,
+    SubBin;
+nif_ioq_payload(heapbin) ->
+    <<"a literal heap binary">>;
+nif_ioq_payload(refcbin) ->
+    iolist_to_binary([lists:seq(1,255) || _ <- lists:seq(1,255)]);
+nif_ioq_payload(Else) ->
+    Else.
 
 %% The NIFs:
 lib_version() -> undefined.
@@ -3032,6 +3208,10 @@ monitor_process_nif(_,_,_,_) -> ?nif_stub.
 demonitor_process_nif(_,_) -> ?nif_stub.
 compare_monitors_nif(_,_) -> ?nif_stub.
 monitor_frenzy_nif(_,_,_,_) -> ?nif_stub.
+ioq_nif(_) -> ?nif_stub.
+ioq_nif(_,_) -> ?nif_stub.
+ioq_nif(_,_,_) -> ?nif_stub.
+ioq_nif(_,_,_,_) -> ?nif_stub.
 
 %% whereis
 whereis_send(_Type,_Name,_Msg) -> ?nif_stub.
diff --git a/erts/emulator/test/nif_SUITE_data/nif_SUITE.c b/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
index 307d1c390f..b47d013bd2 100644
--- a/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
+++ b/erts/emulator/test/nif_SUITE_data/nif_SUITE.c
@@ -186,6 +186,12 @@ static ErlNifResourceTypeInit frenzy_rt_init = {
 
 static ErlNifResourceType* whereis_resource_type;
 static void whereis_thread_resource_dtor(ErlNifEnv* env, void* obj);
+static ErlNifResourceType* ioq_resource_type;
+
+static void ioq_resource_dtor(ErlNifEnv* env, void* obj);
+struct ioq_resource {
+    ErlNifIOQueue *q;
+};
 
 static int get_pointer(ErlNifEnv* env, ERL_NIF_TERM term, void** pp)
 {
@@ -243,6 +249,10 @@ static int load(ErlNifEnv* env, void** priv_data, ERL_NIF_TERM load_info)
     whereis_resource_type = enif_open_resource_type(env, NULL, "nif_SUITE.whereis",
                             whereis_thread_resource_dtor, ERL_NIF_RT_CREATE, NULL);
 
+    ioq_resource_type = enif_open_resource_type(env,NULL,"ioq",
+                                                ioq_resource_dtor,
+                                                ERL_NIF_RT_CREATE, NULL);
+
     atom_false = enif_make_atom(env,"false");
     atom_true = enif_make_atom(env,"true");
     atom_self = enif_make_atom(env,"self");
@@ -2430,7 +2440,6 @@ static ERL_NIF_TERM format_term(ErlNifEnv* env, int argc, const ERL_NIF_TERM arg
     return enif_make_binary(env,&obin);
 }
 
-
 static int get_fd(ErlNifEnv* env, ERL_NIF_TERM term, struct fd_resource** rsrc)
 {
     if (!enif_get_resource(env, term, fd_resource_type, (void**)rsrc)) {
@@ -3158,7 +3167,231 @@ static void frenzy_resource_down(ErlNifEnv* env, void* obj, ErlNifPid* pid,
     abort();
 }
 
+/*********** testing ioq ************/
+
+static void ioq_resource_dtor(ErlNifEnv* env, void* obj) {
+
+}
+
+#ifndef __WIN32__
+static int writeiovec(ErlNifEnv *env, ERL_NIF_TERM term, ERL_NIF_TERM *tail, ErlNifIOQueue *q, int fd) {
+    ErlNifIOVec vec, *iovec = &vec;
+    SysIOVec *sysiovec;
+    int saved_errno;
+    int iovcnt, n;
+
+    if (!enif_inspect_iovec(env, 64, term, tail, &iovec))
+        return -2;
+
+    if (enif_ioq_size(q) > 0) {
+        /* If the I/O queue contains data we enqueue the iovec and then
+           peek the data to write out of the queue. */
+        if (!enif_ioq_enqv(q, iovec, 0))
+            return -3;
+
+        sysiovec = enif_ioq_peek(q, &iovcnt);
+    } else {
+        /* If the I/O queue is empty we skip the trip through it. */
+        iovcnt = iovec->iovcnt;
+        sysiovec = iovec->iov;
+    }
+
+    /* Attempt to write the data */
+    n = writev(fd, sysiovec, iovcnt);
+    saved_errno = errno;
+
+    if (enif_ioq_size(q) == 0) {
+        /* If the I/O queue was initially empty we enqueue any
+           remaining data into the queue for writing later. */
+        if (n >= 0 && !enif_ioq_enqv(q, iovec, n))
+            return -3;
+    } else {
+        /* Dequeue any data that was written from the queue. */
+        if (n > 0 && !enif_ioq_deq(q, n, NULL))
+            return -4;
+    }
+
+    /* return n, which is either number of bytes written or -1 if
+       some error happened */
+    errno = saved_errno;
+    return n;
+}
+#endif
+
+static ERL_NIF_TERM ioq(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[])
+{
+    struct ioq_resource *ioq;
+    ERL_NIF_TERM ret;
+    if (enif_is_identical(argv[0], enif_make_atom(env, "create"))) {
+        ErlNifIOQueue *q = enif_ioq_create(ERL_NIF_IOQ_NORMAL);
+        ioq = (struct ioq_resource *)enif_alloc_resource(ioq_resource_type,
+                                                         sizeof(*ioq));
+        ioq->q = q;
+        ret = enif_make_resource(env, ioq);
+        enif_release_resource(ioq);
+        return ret;
+    } else if (enif_is_identical(argv[0], enif_make_atom(env, "inspect"))) {
+        ErlNifIOVec vec, *iovec = NULL;
+        int i, iovcnt;
+        ERL_NIF_TERM *elems, tail, list;
+        ErlNifEnv *myenv = NULL;
+
+        if (enif_is_identical(argv[2], enif_make_atom(env, "use_stack")))
+            iovec = &vec;
+        if (enif_is_identical(argv[3], enif_make_atom(env, "use_env")))
+            myenv = env;
+        if (!enif_inspect_iovec(myenv, ~(size_t)0, argv[1], &tail, &iovec))
+            return enif_make_badarg(env);
+
+        iovcnt = iovec->iovcnt;
+        elems = enif_alloc(sizeof(ERL_NIF_TERM) * iovcnt);
+
+        for (i = 0; i < iovcnt; i++) {
+            ErlNifBinary bin;
+            if (!enif_alloc_binary(iovec->iov[i].iov_len, &bin)) {
+                enif_free_iovec(iovec);
+                enif_free(elems);
+                return enif_make_badarg(env);
+            }
+            memcpy(bin.data, iovec->iov[i].iov_base, iovec->iov[i].iov_len);
+            elems[i] = enif_make_binary(env, &bin);
+        }
+
+        if (!myenv)
+            enif_free_iovec(iovec);
+
+	list = enif_make_list_from_array(env, elems, iovcnt);
+	enif_free(elems);
+	return list;
+    } else {
+        unsigned skip;
+        if (!enif_get_resource(env, argv[1], ioq_resource_type, (void**)&ioq)
+            || !ioq->q)
+            return enif_make_badarg(env);
+
+        if (enif_is_identical(argv[0], enif_make_atom(env, "example"))) {
+#ifndef __WIN32__
+            int fd[2], res = 0, cnt = 0, queue_cnt;
+            ERL_NIF_TERM tail;
+            char buff[255];
+            pipe(fd);
+            fcntl(fd[0], F_SETFL, fcntl(fd[0], F_GETFL) | O_NONBLOCK);
+            fcntl(fd[1], F_SETFL, fcntl(fd[1], F_GETFL) | O_NONBLOCK);
+
+            /* Write until the pipe buffer is full, which should result in data
+             * being queued up. */
+            for (res = 0; res >= 0; ) {
+                cnt += res;
+                res = writeiovec(env, argv[2], &tail, ioq->q, fd[1]);
+            }
+
+            /* Flush the queue while reading from the other end of the pipe. */
+            tail = enif_make_list(env, 0);
+            while (enif_ioq_size(ioq->q) > 0) {
+                res = writeiovec(env, tail, &tail, ioq->q, fd[1]);
 
+                if (res < 0 && errno != EAGAIN) {
+                    break;
+                } else if (res > 0) {
+                    cnt += res;
+                }
+
+                for (res = 0; res >= 0; ) {
+                    cnt -= res;
+                    res = read(fd[0], buff, sizeof(buff));
+                }
+            }
+
+            close(fd[0]);
+            close(fd[1]);
+
+            /* Check that we read as much as we wrote */
+            if (cnt == 0 && enif_ioq_size(ioq->q) == 0)
+                return enif_make_atom(env, "true");
+
+            return enif_make_int(env, cnt);
+#else
+            return enif_make_atom(env, "true");
+#endif
+        }
+        if (enif_is_identical(argv[0], enif_make_atom(env, "destroy"))) {
+            enif_ioq_destroy(ioq->q);
+            ioq->q = NULL;
+            return enif_make_atom(env, "false");
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "enqv"))) {
+            ErlNifIOVec vec, *iovec = &vec;
+            ERL_NIF_TERM tail;
+
+            if (!enif_get_uint(env, argv[3], &skip))
+                return enif_make_badarg(env);
+            if (!enif_inspect_iovec(env, ~0ul, argv[2], &tail, &iovec))
+                return enif_make_badarg(env);
+            if (!enif_ioq_enqv(ioq->q, iovec, skip))
+                return enif_make_badarg(env);
+
+            return enif_make_atom(env, "true");
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "enqb"))) {
+            ErlNifBinary bin;
+            if (!enif_get_uint(env, argv[3], &skip) ||
+                !enif_inspect_binary(env, argv[2], &bin))
+                return enif_make_badarg(env);
+
+            if (!enif_ioq_enq_binary(ioq->q, &bin, skip))
+                return enif_make_badarg(env);
+
+            return enif_make_atom(env, "true");
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "enqbraw"))) {
+            ErlNifBinary bin;
+            ErlNifBinary localbin;
+	    int i;
+            if (!enif_get_uint(env, argv[3], &skip) ||
+                !enif_inspect_binary(env, argv[2], &bin) ||
+                !enif_alloc_binary(bin.size, &localbin))
+                return enif_make_badarg(env);
+
+            memcpy(localbin.data, bin.data, bin.size);
+            i = enif_ioq_enq_binary(ioq->q, &localbin, skip);
+	    if (!i)
+		return enif_make_badarg(env);
+	    else
+		return enif_make_atom(env, "true");
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "peek"))) {
+            int iovlen, num, i, off = 0;
+            SysIOVec *iov = enif_ioq_peek(ioq->q, &iovlen);
+            ErlNifBinary bin;
+
+            if (!enif_get_int(env, argv[2], &num) || !enif_alloc_binary(num, &bin))
+                return enif_make_badarg(env);
+
+            for (i = 0; i < iovlen && num > 0; i++) {
+                int to_copy = num < iov[i].iov_len ? num : iov[i].iov_len;
+                memcpy(bin.data + off, iov[i].iov_base, to_copy);
+                num -= to_copy;
+                off += to_copy;
+            }
+
+            return enif_make_binary(env, &bin);
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "deq"))) {
+            int num;
+            size_t sz;
+            ErlNifUInt64 sz64;
+            if (!enif_get_int(env, argv[2], &num))
+                return enif_make_badarg(env);
+
+            if (!enif_ioq_deq(ioq->q, num, &sz))
+                return enif_make_badarg(env);
+
+            sz64 = sz;
+
+            return enif_make_uint64(env, sz64);
+        } else if (enif_is_identical(argv[0], enif_make_atom(env, "size"))) {
+            ErlNifUInt64 size = enif_ioq_size(ioq->q);
+            return enif_make_uint64(env, size);
+        }
+    }
+
+    return enif_make_badarg(env);
+}
 
 static ErlNifFunc nif_funcs[] =
 {
@@ -3255,7 +3488,11 @@ static ErlNifFunc nif_funcs[] =
     {"whereis_send", 3, whereis_send},
     {"whereis_term", 2, whereis_term},
     {"whereis_thd_lookup", 2, whereis_thd_lookup},
-    {"whereis_thd_result", 1, whereis_thd_result}
+    {"whereis_thd_result", 1, whereis_thd_result},
+    {"ioq_nif", 1, ioq},
+    {"ioq_nif", 2, ioq},
+    {"ioq_nif", 3, ioq},
+    {"ioq_nif", 4, ioq}
 };
 
 ERL_NIF_INIT(nif_SUITE,nif_funcs,load,NULL,upgrade,unload)