4 files changed, 25 insertions, 31 deletions

diff --git a/src/packet_stream.c b/src/packet_stream.c
index 110984b..bd6d476 100644
--- a/src/packet_stream.c
+++ b/src/packet_stream.c
@@ -3,7 +3,6 @@
 enum {
     PACKET_STREAM_CONNECTING = 0,
     PACKET_STREAM_CONNECTED,
-    PACKET_STREAM_FLUSHING,
     PACKET_STREAM_DISCONNECTING,
     PACKET_STREAM_DISCONNECTED
 };
@@ -153,13 +152,6 @@ static void stream_write_callback(struct ev_loop *loop, ev_io *w, s32 revents)
             stream->out.index = 0;
         } else {
             stop_write_internal(stream);
-            if (stream->connect == PACKET_STREAM_FLUSHING) {
-                if (stream->corked) {
-                    stop_internal(stream);
-                } else {
-                    shutdown_internal(stream);
-                }
-            }
             return;
         }
     }
@@ -243,8 +235,7 @@ void aki_packet_stream_cork(struct aki_packet_stream *stream, bool cork)
 
 bool aki_packet_stream_send_packet(struct aki_packet_stream *stream, struct aki_packet *packet)
 {
-    if (stream->connect == PACKET_STREAM_FLUSHING ||
-        stream->connect == PACKET_STREAM_DISCONNECTING) {
+    if (stream->connect == PACKET_STREAM_DISCONNECTING) {
         return false;
     }
     al_assert(stream->connect == PACKET_STREAM_CONNECTED);
@@ -256,22 +247,6 @@ bool aki_packet_stream_send_packet(struct aki_packet_stream *stream, struct aki_
     return true;
 }
 
-void aki_packet_stream_flush(struct aki_packet_stream *stream)
-{
-    al_assert(stream->connect == PACKET_STREAM_CONNECTED ||
-        stream->connect == PACKET_STREAM_CONNECTING);
-    // If connect = CONNECTING, we will wait utill connected then flush.
-    if (!stream->out.active) {
-        if (stream->corked) {
-            stop_internal(stream);
-        } else {
-            shutdown_internal(stream);
-        }
-    } else {
-        stream->connect = PACKET_STREAM_FLUSHING;
-    }
-}
-
 void aki_packet_stream_disconnect(struct aki_packet_stream *stream)
 {
     al_assert(stream->connect != PACKET_STREAM_DISCONNECTING);
diff --git a/src/packet_stream.h b/src/packet_stream.h
index 180c507..a6ca708 100644
--- a/src/packet_stream.h
+++ b/src/packet_stream.h
@@ -43,6 +43,5 @@ void aki_packet_stream_connect(struct aki_packet_stream *stream, struct aki_even
 void aki_packet_stream_reconnect(struct aki_packet_stream *stream, str *addr, u16 port);
 void aki_packet_stream_cork(struct aki_packet_stream *stream, bool cork);
 bool aki_packet_stream_send_packet(struct aki_packet_stream *stream, struct aki_packet *packet);
-void aki_packet_stream_flush(struct aki_packet_stream *stream);
 void aki_packet_stream_disconnect(struct aki_packet_stream *stream);
 void aki_packet_stream_free(struct aki_packet_stream *stream);
diff --git a/src/rpc.c b/src/rpc.c
index e086d89..b9f74c5 100644
--- a/src/rpc.c
+++ b/src/rpc.c
@@ -42,6 +42,7 @@ static void packet_callback(void *userdata, struct aki_packet_stream *stream, st
                 aki_packet_write_s8(rpacket, -1);
                 aki_packet_write_u32(rpacket, id);
                 if (command->callback(command->userdata, conn, packet, rpacket)) {
+                    conn->outgoing++;
                     aki_packet_stream_send_packet(stream, rpacket);
                 } else {
                     aki_packet_free(rpacket);
@@ -55,8 +56,13 @@ static void packet_callback(void *userdata, struct aki_packet_stream *stream, st
 
 static void packet_sent_callback(void *userdata, struct aki_packet *packet)
 {
-    (void)userdata;
+    struct aki_rpc_connection *conn = (struct aki_rpc_connection *)userdata;
     aki_packet_free(packet);
+    al_assert(conn->outgoing > 0);
+    conn->outgoing--;
+    if (conn->flushing && conn->outgoing == 0) {
+        aki_packet_stream_disconnect(&conn->stream);
+    }
 }
 
 static void stream_connection_callback(void *userdata, struct aki_packet_stream *stream)
@@ -81,11 +87,18 @@ static void stream_connection_closed_callback(void *userdata, struct aki_packet_
     conn->rpc->connection_closed_callback(conn->rpc->userdata, conn);
 }
 
+static void init_connection(struct aki_rpc_connection *conn)
+{
+    al_array_init(conn->callbacks);
+    conn->outgoing = 0;
+    conn->flushing = false;
+}
+
 void aki_rpc_add_socket(struct aki_rpc *rpc, struct aki_socket *sock)
 {
     struct aki_rpc_connection *conn = al_alloc_object(struct aki_rpc_connection);
     conn->rpc = rpc;
-    al_array_init(conn->callbacks);
+    init_connection(conn);
     conn->stream.connection_callback = stream_connection_callback;
     conn->stream.connection_closed_callback = stream_connection_closed_callback;
     conn->stream.userdata = conn;
@@ -103,7 +116,7 @@ bool aki_rpc_prepare_client(struct aki_rpc *rpc, u8 type, u8 id)
         return false;
     }
     conn->rpc = rpc;
-    al_array_init(conn->callbacks);
+    init_connection(conn);
     if (conn->stream.sock.type == AKI_SOCKET_TCP) {
         aki_packet_stream_set_nodelay(&conn->stream, 1);
     }
@@ -154,12 +167,17 @@ void aki_rpc_connection_command(struct aki_rpc_connection *conn, struct aki_pack
             .userdata = userdata
         }));
     }
+    conn->outgoing++;
     aki_packet_stream_send_packet(&conn->stream, packet);
 }
 
 void aki_rpc_conn_flush(struct aki_rpc_connection *conn)
 {
-    aki_packet_stream_flush(&conn->stream);
+    if (conn->outgoing == 0) {
+        aki_packet_stream_disconnect(&conn->stream);
+    } else {
+        conn->flushing = true;
+    }
 }
 
 void aki_rpc_conn_disconnect(struct aki_rpc_connection *conn)
diff --git a/src/rpc2.h b/src/rpc2.h
index d03a8a7..f96a1e5 100644
--- a/src/rpc2.h
+++ b/src/rpc2.h
@@ -15,6 +15,8 @@ struct aki_rpc_callback {
 struct aki_rpc_connection {
     struct aki_packet_stream stream;
     array(struct aki_rpc_callback) callbacks;
+    u32 outgoing;
+    bool flushing;
     struct aki_rpc *rpc;
 };