in_tcp: implement pause and resume callbacks

Signed-off-by: Eduardo Silva <[email protected]>

Author: edsiper

plugins/in_tcp/tcp.c

@@ -42,9 +42,8 @@ static int in_tcp_collect(struct flb_input_instance *in,
     connection = flb_downstream_conn_get(ctx->downstream);
 
     if (connection == NULL) {
-        flb_plg_error(ctx->ins, "could not accept new connection");
-
-        return -1;
+        /* connection dropped (e.g. paused) */
+        return 0;
     }
 
     flb_plg_trace(ctx->ins, "new TCP connection arrived FD=%i", connection->fd);
@@ -145,6 +144,37 @@ static int in_tcp_exit(void *data, struct flb_config *config)
     return 0;
 }
 
+static void in_tcp_pause(void *data, struct flb_config *config)
+{
+    struct flb_in_tcp_config *ctx = data;
+    struct mk_list *head;
+    struct mk_list *tmp;
+    struct tcp_conn *conn;
+
+    (void) config;
+
+    flb_downstream_pause(ctx->downstream);
+
+    mk_list_foreach_safe(head, tmp, &ctx->connections) {
+        conn = mk_list_entry(head, struct tcp_conn, _head);
+        if (conn->busy) {
+            conn->pending_close = FLB_TRUE;
+            continue;
+        }
+
+        tcp_conn_del(conn);
+    }
+}
+
+static void in_tcp_resume(void *data, struct flb_config *config)
+{
+    struct flb_in_tcp_config *ctx = data;
+
+    (void) config;
+
+    flb_downstream_resume(ctx->downstream);
+}
+
 static struct flb_config_map config_map[] = {
     {
      FLB_CONFIG_MAP_STR, "format", (char *)NULL,
@@ -183,6 +213,8 @@ struct flb_input_plugin in_tcp_plugin = {
     .cb_pre_run   = NULL,
     .cb_collect   = in_tcp_collect,
     .cb_flush_buf = NULL,
+    .cb_pause     = in_tcp_pause,
+    .cb_resume    = in_tcp_resume,
     .cb_exit      = in_tcp_exit,
     .config_map   = config_map,
     .flags        = FLB_INPUT_NET_SERVER | FLB_IO_OPT_TLS

plugins/in_tcp/tcp_conn.c

@@ -260,6 +260,7 @@ int tcp_conn_event(void *data)
     struct tcp_conn *conn;
     struct flb_connection *connection;
     struct flb_in_tcp_config *ctx;
+    int ret = 0;
 
     connection = (struct flb_connection *) data;
 
@@ -269,20 +270,24 @@ int tcp_conn_event(void *data)
 
     event = &connection->event;
 
+    conn->busy = FLB_TRUE;
+
     if (event->mask & MK_EVENT_READ) {
         available = (conn->buf_size - conn->buf_len) - 1;
         if (available < 1) {
             if (conn->buf_size + ctx->chunk_size > ctx->buffer_size) {
                 flb_plg_warn(ctx->ins,
                              "fd=%i incoming data exceeds 'Buffer_Size' (%zu KB)",
                              event->fd, (ctx->buffer_size / 1024));
+                conn->busy = FLB_FALSE;
                 tcp_conn_del(conn);
                 return -1;
             }
 
             size = conn->buf_size + ctx->chunk_size;
             tmp = flb_realloc(conn->buf_data, size);
             if (!tmp) {
+                conn->busy = FLB_FALSE;
                 flb_errno();
                 return -1;
             }
@@ -301,6 +306,7 @@ int tcp_conn_event(void *data)
 
         if (bytes <= 0) {
             flb_plg_trace(ctx->ins, "fd=%i closed connection", event->fd);
+            conn->busy = FLB_FALSE;
             tcp_conn_del(conn);
             return -1;
         }
@@ -325,23 +331,27 @@ int tcp_conn_event(void *data)
             ret_payload = parse_payload_json(conn);
             if (ret_payload == 0) {
                 /* Incomplete JSON message, we need more data */
-                return -1;
+                ret = -1;
+                goto cleanup;
             }
             else if (ret_payload == -1) {
                 flb_pack_state_reset(&conn->pack_state);
                 flb_pack_state_init(&conn->pack_state);
                 conn->pack_state.multiple = FLB_TRUE;
-                return -1;
+                ret = -1;
+                goto cleanup;
             }
         }
         else if (ctx->format == FLB_TCP_FMT_NONE) {
             ret_payload = parse_payload_none(conn);
             if (ret_payload == 0) {
-                return -1;
+                ret = -1;
+                goto cleanup;
             }
             else if (ret_payload == -1) {
                 conn->buf_len = 0;
-                return -1;
+                ret = -1;
+                goto cleanup;
             }
         }
 
@@ -357,16 +367,27 @@ int tcp_conn_event(void *data)
             conn->pack_state.buf_len = 0;
         }
 
-        return bytes;
+        ret = bytes;
+        goto cleanup;
     }
 
     if (event->mask & MK_EVENT_CLOSE) {
         flb_plg_trace(ctx->ins, "fd=%i hangup", event->fd);
+        conn->busy = FLB_FALSE;
         tcp_conn_del(conn);
         return -1;
     }
 
-    return 0;
+    ret = 0;
+
+cleanup:
+    conn->busy = FLB_FALSE;
+    if (conn->pending_close) {
+        tcp_conn_del(conn);
+        return -1;
+    }
+
+    return ret;
 }
 
 /* Create a new mqtt request instance */
@@ -432,6 +453,9 @@ struct tcp_conn *tcp_conn_add(struct flb_connection *connection,
 
     mk_list_add(&conn->_head, &ctx->connections);
 
+    conn->busy = FLB_FALSE;
+    conn->pending_close = FLB_FALSE;
+
     return conn;
 }
 

plugins/in_tcp/tcp_conn.h

@@ -50,6 +50,9 @@ struct tcp_conn {
     struct flb_pack_state pack_state; /* Internal JSON parser              */
     struct flb_connection *connection;
 
+    int busy;                         /* Connection is being processed     */
+    int pending_close;                /* Defer closing until processing ends */
+
     struct mk_list _head;
 };