Adding upstream version 3.2.6.upstream/3.2.6upstream

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 243 insertions, 0 deletions

diff --git a/src/contrib/conn_pool.c b/src/contrib/conn_pool.c
new file mode 100644
index 0000000..8367555
--- /dev/null
+++ b/src/contrib/conn_pool.c
@@ -0,0 +1,243 @@
+/*  Copyright (C) 2022 CZ.NIC, z.s.p.o. <knot-dns@labs.nic.cz>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+#include "contrib/conn_pool.h"
+
+#include "contrib/sockaddr.h"
+
+conn_pool_t *global_conn_pool = NULL;
+
+static int pool_pop(conn_pool_t *pool, size_t i);
+
+/*!
+ * \brief Try to get an open connection older than specified timestamp.
+ *
+ * \param pool           Pool to search in.
+ * \param older_than     Timestamp that the connection must be older than.
+ * \param next_oldest    Out: the timestamp of the oldest connection (other than the returned).
+ *
+ * \return -1 if error (no such connection), >= 0 connection file descriptor.
+ *
+ * \warning The returned connection is not necessarily the oldest one.
+ */
+static int get_old(conn_pool_t *pool,
+                   knot_time_t older_than,
+                   knot_time_t *next_oldest)
+{
+	assert(pool);
+
+	*next_oldest = 0;
+
+	int fd = -1;
+	pthread_mutex_lock(&pool->mutex);
+
+	for (size_t i = 0; i < pool->capacity; i++) {
+		knot_time_t la = pool->conns[i].last_active;
+		if (fd == -1 && knot_time_cmp(la, older_than) < 0) {
+			fd = pool_pop(pool, i);
+		} else if (knot_time_cmp(la, *next_oldest) < 0) {
+			*next_oldest = la;
+		}
+	}
+
+	pthread_mutex_unlock(&pool->mutex);
+	return fd;
+}
+
+static void *closing_thread(void *_arg)
+{
+	conn_pool_t *pool = _arg;
+
+	while (true) {
+		knot_time_t now = knot_time(), next = 0;
+		knot_timediff_t timeout = conn_pool_timeout(pool, 0);
+		assert(timeout != 0);
+
+		while (true) {
+			int old_fd = get_old(pool, now - timeout + 1, &next);
+			if (old_fd >= 0) {
+				close(old_fd);
+			} else {
+				break;
+			}
+		}
+
+		if (next == 0) {
+			sleep(timeout);
+		} else {
+			sleep(next + timeout - now);
+		}
+	}
+
+	return NULL; // we never get here since the thread will be cancelled instead
+}
+
+conn_pool_t *conn_pool_init(size_t capacity, knot_timediff_t timeout)
+{
+	if (capacity == 0 || timeout == 0) {
+		return NULL;
+	}
+
+	conn_pool_t *pool = calloc(1, sizeof(*pool) + capacity * sizeof(pool->conns[0]));
+	if (pool != NULL) {
+		pool->capacity = capacity;
+		pool->timeout = timeout;
+		if (pthread_mutex_init(&pool->mutex, 0) != 0) {
+			free(pool);
+			return NULL;
+		}
+		if (pthread_create(&pool->closing_thread, NULL, closing_thread, pool) != 0) {
+			pthread_mutex_destroy(&pool->mutex);
+			free(pool);
+			return NULL;
+		}
+	}
+	return pool;
+}
+
+void conn_pool_deinit(conn_pool_t *pool)
+{
+	if (pool != NULL) {
+		pthread_cancel(pool->closing_thread);
+		pthread_join(pool->closing_thread, NULL);
+
+		int fd;
+		knot_time_t unused;
+		while ((fd = get_old(pool, 0, &unused)) >= 0) {
+			close(fd);
+		}
+
+		pthread_mutex_destroy(&pool->mutex);
+		free(pool);
+	}
+}
+
+knot_timediff_t conn_pool_timeout(conn_pool_t *pool,
+                                  knot_timediff_t new_timeout)
+{
+	if (pool == NULL) {
+		return 0;
+	}
+
+	pthread_mutex_lock(&pool->mutex);
+
+	knot_timediff_t prev = pool->timeout;
+	if (new_timeout != 0) {
+		pool->timeout = new_timeout;
+	}
+
+	pthread_mutex_unlock(&pool->mutex);
+	return prev;
+}
+
+static int pool_pop(conn_pool_t *pool, size_t i)
+{
+	conn_pool_memb_t *conn = &pool->conns[i];
+	assert(conn->last_active != 0);
+	assert(pool->usage > 0);
+	int fd = conn->fd;
+	memset(conn, 0, sizeof(*conn));
+	pool->usage--;
+	return fd;
+}
+
+int conn_pool_get(conn_pool_t *pool,
+                  struct sockaddr_storage *src,
+                  struct sockaddr_storage *dst)
+{
+	if (pool == NULL) {
+		return -1;
+	}
+
+	int fd = -1;
+	pthread_mutex_lock(&pool->mutex);
+
+	for (size_t i = 0; i < pool->capacity; i++) {
+		if (pool->conns[i].last_active != 0 &&
+		    sockaddr_cmp(&pool->conns[i].dst, dst, false) == 0 &&
+		    sockaddr_cmp(&pool->conns[i].src, src, true) == 0) {
+			fd = pool_pop(pool, i);
+			break;
+		}
+	}
+
+	pthread_mutex_unlock(&pool->mutex);
+
+	if (fd >= 0) {
+		uint8_t unused;
+		int peek = recv(fd, &unused, 1, MSG_PEEK | MSG_DONTWAIT);
+		if (peek >= 0) { // closed or pending data
+			close(fd);
+			fd = -1;
+		}
+	}
+
+	return fd;
+}
+
+static void pool_push(conn_pool_t *pool, size_t i,
+                      struct sockaddr_storage *src,
+                      struct sockaddr_storage *dst,
+                      int fd)
+{
+	conn_pool_memb_t *conn = &pool->conns[i];
+	assert(conn->last_active == 0);
+	assert(pool->usage < pool->capacity);
+	conn->last_active = knot_time();
+	conn->fd = fd;
+	memcpy(&conn->src, src, sizeof(conn->src));
+	memcpy(&conn->dst, dst, sizeof(conn->dst));
+	pool->usage++;
+}
+
+int conn_pool_put(conn_pool_t *pool,
+                  struct sockaddr_storage *src,
+                  struct sockaddr_storage *dst,
+                  int fd)
+{
+	if (pool == NULL || pool->capacity == 0) {
+		return fd;
+	}
+
+	knot_time_t oldest_time = 0;
+	size_t oldest_i = pool->capacity;
+
+	pthread_mutex_lock(&pool->mutex);
+
+	for (size_t i = 0; i < pool->capacity; i++) {
+		knot_time_t la = pool->conns[i].last_active;
+		if (la == 0) {
+			pool_push(pool, i, src, dst, fd);
+			pthread_mutex_unlock(&pool->mutex);
+			return -1;
+		} else if (knot_time_cmp(la, oldest_time) < 0) {
+			oldest_time = la;
+			oldest_i = i;
+		}
+	}
+
+	assert(oldest_i < pool->capacity);
+	int oldest_fd = pool_pop(pool, oldest_i);
+	pool_push(pool, oldest_i, src, dst, fd);
+	pthread_mutex_unlock(&pool->mutex);
+	return oldest_fd;
+}