From 3f1ef40b9011276eb8bdc366c5ef1e25d79befa5 Mon Sep 17 00:00:00 2001
From: Maas-Maarten Zeeman <mmzeeman@xs4all.nl>
Date: Sun, 26 Feb 2017 14:55:47 +0100
Subject: [PATCH] Fix reference count to connection when a prepared statement
 is made

---
 c_src/esqlite3_nif.c  |  9 ++-------
 src/esqlite3.erl      | 21 ++++++++++++---------
 test/esqlite_test.erl | 22 +++++++++++++++++++++-
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/c_src/esqlite3_nif.c b/c_src/esqlite3_nif.c
index d79d3f9..773fcaf 100644
--- a/c_src/esqlite3_nif.c
+++ b/c_src/esqlite3_nif.c
@@ -242,7 +242,7 @@ destruct_esqlite_statement(ErlNifEnv *env, void *arg)
 	   stmt->statement = NULL;
     }
 
-    enif_release_resource(stmt->connection);
+    stmt->connection = NULL;
 }
 
 static ERL_NIF_TERM
@@ -681,7 +681,7 @@ esqlite_connection_run(void *arg)
     
 	    command_destroy(cmd);    
     }
-  
+
     return NULL;
 }
 
@@ -869,11 +869,6 @@ esqlite_prepare(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
     if(!cmd) 
 	    return make_error_tuple(env, "command_create_failed");
 
-    /* Keep a reference to the connection to prevent it from being taken down
-     * while the prepare statement is waiting on the queue.
-     */
-    enif_keep_resource(conn);
-
     cmd->type = cmd_prepare;
     cmd->ref = enif_make_copy(cmd->env, argv[1]);
     cmd->pid = pid;
diff --git a/src/esqlite3.erl b/src/esqlite3.erl
index 3f5b1b2..17a9ef2 100644
--- a/src/esqlite3.erl
+++ b/src/esqlite3.erl
@@ -40,8 +40,8 @@
 -define(DEFAULT_TIMEOUT, 5000).
 
 %%
--type connection() :: tuple().
--type statement() :: term().
+-type connection() :: {connection, reference(), term()}.
+-type statement() :: {statement, term(), connection()}.
 -type sql() :: iolist().
 
 %% @doc Opens a sqlite3 database mentioned in Filename.
@@ -260,10 +260,13 @@ prepare(Sql, Connection) ->
 %% @doc
 %%
 %% @spec(iolist(), connection(), timeout()) -> {ok, prepared_statement()} | {error, error_message()}
-prepare(Sql, {connection, _Ref, Connection}, Timeout) ->
+prepare(Sql, {connection, _Ref, Connection}=C, Timeout) ->
     Ref = make_ref(),
     ok = esqlite3_nif:prepare(Connection, Ref, self(), Sql),
-    receive_answer(Ref, Timeout).
+    case receive_answer(Ref, Timeout) of
+        {ok, Stmt} -> {ok, {statement, Stmt, C}};
+        Else -> Else
+    end.
 
 %% @doc Step
 %%
@@ -275,7 +278,7 @@ step(Stmt) ->
 %%
 %% @spec step(prepared_statement(), timeout()) -> tuple()
 -spec step(term(), timeout()) -> tuple() | '$busy' | '$done'.
-step(Stmt, Timeout) ->
+step({statement, Stmt, _}, Timeout) ->
     Ref = make_ref(),
     ok = esqlite3_nif:step(Stmt, Ref, self()),
     receive_answer(Ref, Timeout).
@@ -283,7 +286,7 @@ step(Stmt, Timeout) ->
 %% @doc Reset the prepared statement back to its initial state.
 %%
 %% @spec reset(prepared_statement()) -> ok | {error, error_message()}
-reset(Stmt) ->
+reset({statement, Stmt, _}) ->
     Ref = make_ref(),
     ok = esqlite3_nif:reset(Stmt, Ref, self()),
     receive_answer(Ref, ?DEFAULT_TIMEOUT).
@@ -297,7 +300,7 @@ bind(Stmt, Args) ->
 %% @doc Bind values to prepared statements
 %%
 %% @spec bind(prepared_statement(), [], timeout()) -> ok | {error, error_message()}
-bind(Stmt, Args, Timeout) ->
+bind({statement, Stmt, _}, Args, Timeout) ->
     Ref = make_ref(),
     ok = esqlite3_nif:bind(Stmt, Ref, self(), Args),
     receive_answer(Ref, Timeout).
@@ -309,7 +312,7 @@ column_names(Stmt) ->
     column_names(Stmt, ?DEFAULT_TIMEOUT).
 
 -spec column_names(statement(), timeout()) -> {atom()}.
-column_names(Stmt, Timeout) ->
+column_names({statement, Stmt, _}, Timeout) ->
     Ref = make_ref(),
     ok = esqlite3_nif:column_names(Stmt, Ref, self()),
     receive_answer(Ref, Timeout).
@@ -321,7 +324,7 @@ column_types(Stmt) ->
     column_types(Stmt, ?DEFAULT_TIMEOUT).
 
 -spec column_types(statement(), timeout()) -> {atom()}.
-column_types(Stmt, Timeout) ->
+column_types({statement, Stmt, _}, Timeout) ->
     Ref = make_ref(),
     ok = esqlite3_nif:column_types(Stmt, Ref, self()),
     receive_answer(Ref, Timeout).
diff --git a/test/esqlite_test.erl b/test/esqlite_test.erl
index 4247739..233c456 100644
--- a/test/esqlite_test.erl
+++ b/test/esqlite_test.erl
@@ -115,7 +115,6 @@ bind_test() ->
     %% utf-8
     ?assertEqual([{<<228,184,138,230,181,183>>, 100}], 
         esqlite3:q("select one, two from test_table where two = 100", Db)),
-
     
     ok.
 
@@ -326,3 +325,24 @@ sqlite_source_id_test() ->
     ?assertEqual({row, {<<"2015-07-23 16:39:33 793e206f9032d9205bdb3f447b136bed9a25fa22">>}}, esqlite3:step(Stmt)),
     ok.
     
+garbage_collect_test() ->
+    F = fun() ->
+        {ok, Db} = esqlite3:open(":memory:"),
+        [] = esqlite3:q("create table test(one, two, three)", Db),
+        {ok, Stmt} = esqlite3:prepare("select * from test", Db),
+        '$done' = esqlite3:step(Stmt)
+    end,
+
+    [spawn(F) || _X <- lists:seq(0,30)],
+    receive after 500 -> ok end,
+    erlang:garbage_collect(),
+
+    [spawn(F) || _X <- lists:seq(0,30)],
+    receive after 500 -> ok end,
+    erlang:garbage_collect(),
+
+    ok.
+
+
+
+        
\ No newline at end of file