reimplementation of the cache networking protocol to be fully non-blocking

core/buffer.c

@@ -319,3 +319,12 @@ int uwsgi_buffer_send(struct uwsgi_buffer *ub, int fd) {
 
 	return 0;
 }
+
+int uwsgi_buffer_set_uh(struct uwsgi_buffer *ub, uint8_t modifier1, uint8_t modifier2) {
+	if (ub->pos < 4) return -1;
+	ub->buf[0] = modifier1;
+	ub->buf[1] = (uint8_t) ((ub->pos - 4) & 0xff);
+	ub->buf[2] = (uint8_t) (((ub->pos - 4) >> 8) & 0xff);
+	ub->buf[3] = modifier2;
+	return 0;
+}

core/cache.c

@@ -1095,6 +1095,20 @@ struct uwsgi_cache *uwsgi_cache_by_name(char *name) {
 	return NULL;
 }
 
+struct uwsgi_cache *uwsgi_cache_by_namelen(char *name, uint16_t len) {
+        struct uwsgi_cache *uc = uwsgi.caches;
+        if (!name || *name == 0) {
+                return uwsgi.caches;
+        }
+        while(uc) {
+                if (uc->name && !uwsgi_strncmp(uc->name, uc->name_len, name, len)) {
+                        return uc;
+                }
+                uc = uc->next;
+        }
+        return NULL;
+}
+
 void uwsgi_cache_create_all() {
 
 	if (uwsgi.cache_setup) return;
@@ -1125,15 +1139,150 @@ void uwsgi_cache_create_all() {
  * You have to free the returned memory !!!
  *
  */
-char *uwsgi_cache_magic_get(char *key, uint16_t keylen, uint64_t *vallen, char *cachename) {
+
+void uwsgi_cache_magic_context_hook(char *key, uint16_t key_len, char *value, uint16_t vallen, void *data) {
+	struct uwsgi_cache_magic_context *ucmc = (struct uwsgi_cache_magic_context *) data;
+
+	if (!uwsgi_strncmp(key, key_len, "cmd", 3)) {
+		ucmc->cmd = value;
+		ucmc->cmd_len = vallen;
+		return;
+	}
+
+	if (!uwsgi_strncmp(key, key_len, "key", 3)) {
+                ucmc->key = value;
+                ucmc->key_len = vallen; 
+                return;
+        }
+
+	if (!uwsgi_strncmp(key, key_len, "expires", 7)) {
+                ucmc->expires = uwsgi_str_num(value, vallen);
+                return;
+        }
+
+	if (!uwsgi_strncmp(key, key_len, "size", 4)) {
+                ucmc->size = uwsgi_str_num(value, vallen);
+                return;
+        }
+
+	if (!uwsgi_strncmp(key, key_len, "cache", 5)) {
+                ucmc->cache = value;
+                ucmc->cache_len = vallen; 
+                return;
+        }
+
+	if (!uwsgi_strncmp(key, key_len, "status", 6)) {
+                ucmc->status = value;
+                ucmc->status_len = vallen;
+                return;
+        }
+}
+
+static struct uwsgi_buffer *uwsgi_cache_prepare_magic_get(char *cache_name, uint16_t cache_name_len, char *key, uint16_t key_len) {
+	struct uwsgi_buffer *ub = uwsgi_buffer_new(uwsgi.page_size);
+	ub->pos = 4;
+
+	if (uwsgi_buffer_append_keyval(ub, "cmd", 3, "get", 3)) goto error;
+	if (uwsgi_buffer_append_keyval(ub, "key", 3, key, key_len)) goto error;
+	if (cache_name) {
+		if (uwsgi_buffer_append_keyval(ub, "cache", 5, cache_name, cache_name_len)) goto error;
+	}
+
+	return ub;
+error:
+	uwsgi_buffer_destroy(ub);
+	return NULL;
+}
+
+struct uwsgi_buffer *uwsgi_cache_prepare_magic_exists(char *cache_name, uint16_t cache_name_len, char *key, uint16_t key_len) {
+        struct uwsgi_buffer *ub = uwsgi_buffer_new(uwsgi.page_size);
+	ub->pos = 4;
+
+        if (uwsgi_buffer_append_keyval(ub, "cmd", 3, "exists", 6)) goto error;
+        if (uwsgi_buffer_append_keyval(ub, "key", 3, key, key_len)) goto error;
+        if (cache_name) {
+                if (uwsgi_buffer_append_keyval(ub, "cache", 5, cache_name, cache_name_len)) goto error;
+        }
+
+        return ub;
+error:
+        uwsgi_buffer_destroy(ub);
+        return NULL;
+}
+
+struct uwsgi_buffer *uwsgi_cache_prepare_magic_del(char *cache_name, uint16_t cache_name_len, char *key, uint16_t key_len) {
+        struct uwsgi_buffer *ub = uwsgi_buffer_new(uwsgi.page_size);
+	ub->pos = 4;
+
+        if (uwsgi_buffer_append_keyval(ub, "cmd", 3, "del", 3)) goto error;
+        if (uwsgi_buffer_append_keyval(ub, "key", 3, key, key_len)) goto error;
+        if (cache_name) {
+                if (uwsgi_buffer_append_keyval(ub, "cache", 5, cache_name, cache_name_len)) goto error;
+        }
+
+        return ub;
+error:
+        uwsgi_buffer_destroy(ub);
+        return NULL;
+}
+
+struct uwsgi_buffer *uwsgi_cache_prepare_magic_set(char *cache_name, uint16_t cache_name_len, char *key, uint16_t key_len, uint64_t len, uint64_t expires) {
+        struct uwsgi_buffer *ub = uwsgi_buffer_new(uwsgi.page_size);
+	ub->pos = 4;
+
+        if (uwsgi_buffer_append_keyval(ub, "cmd", 3, "set", 3)) goto error;
+        if (uwsgi_buffer_append_keyval(ub, "key", 3, key, key_len)) goto error;
+        if (uwsgi_buffer_append_keynum(ub, "size", 4, len)) goto error;
+	if (expires > 0) {
+		if (uwsgi_buffer_append_keynum(ub, "expires", 7, expires)) goto error;
+	}
+        if (uwsgi_buffer_append_keynum(ub, "size", 4, len)) goto error;
+        if (cache_name) {
+                if (uwsgi_buffer_append_keyval(ub, "cache", 5, cache_name, cache_name_len)) goto error;
+        }
+
+        return ub;
+error:
+        uwsgi_buffer_destroy(ub);
+        return NULL;
+}
+
+struct uwsgi_buffer *uwsgi_cache_prepare_magic_update(char *cache_name, uint16_t cache_name_len, char *key, uint16_t key_len, uint64_t len, uint64_t expires) {
+        struct uwsgi_buffer *ub = uwsgi_buffer_new(uwsgi.page_size);
+	ub->pos = 4;
+
+        if (uwsgi_buffer_append_keyval(ub, "cmd", 3, "update", 6)) goto error;
+        if (uwsgi_buffer_append_keyval(ub, "key", 3, key, key_len)) goto error;
+        if (uwsgi_buffer_append_keynum(ub, "size", 4, len)) goto error;
+        if (expires > 0) {
+                if (uwsgi_buffer_append_keynum(ub, "expires", 7, expires)) goto error;
+        }
+        if (uwsgi_buffer_append_keynum(ub, "size", 4, len)) goto error;
+        if (cache_name) {
+                if (uwsgi_buffer_append_keyval(ub, "cache", 5, cache_name, cache_name_len)) goto error;
+        }
+
+        return ub;
+error:
+        uwsgi_buffer_destroy(ub);
+        return NULL;
+}
+
+char *uwsgi_cache_magic_get(char *key, uint16_t keylen, uint64_t *vallen, char *cache) {
+	struct uwsgi_cache_magic_context ucmc;
 	struct uwsgi_cache *uc = NULL;
 	char *cache_server = NULL;
-	if (cachename) {
-		char *at = strchr(cachename, '@');
+	char *cache_name = NULL;
+	uint16_t cache_name_len = 0;
+	if (cache) {
+		char *at = strchr(cache, '@');
 		if (!at) {
-			uc = uwsgi_cache_by_name(cachename);
+			uc = uwsgi_cache_by_name(cache);
 		}
 		else {
+			cache_server = at + 1;
+			cache_name = cache;
+			cache_name_len = at - cache;
 		}
 	}
 	// use default (local) cache
@@ -1157,6 +1306,88 @@ char *uwsgi_cache_magic_get(char *key, uint16_t keylen, uint64_t *vallen, char *
 
 	// we have a remote one
 	if (cache_server) {
+		int fd = uwsgi_connect(cache_server, 0, 1);
+		if (fd < 0) return NULL;
+
+		// wait for connection;
+		int ret = uwsgi.wait_write_hook(fd, uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT]);
+		if (ret <= 0) {
+			close(fd);
+			return NULL;
+		}
+
+		struct uwsgi_buffer *ub = uwsgi_cache_prepare_magic_get(cache_name, cache_name_len, key, keylen);
+		if (!ub) {
+			close(fd);
+			return NULL;
+		}
+
+		if (uwsgi_buffer_set_uh(ub, 111, 17)) {
+			close(fd);
+			uwsgi_buffer_destroy(ub);
+			return NULL;
+		}
+
+		if (uwsgi_write_true_nb(fd, ub->buf, ub->pos, uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT])) {
+			close(fd);
+			uwsgi_buffer_destroy(ub);
+			return NULL;
+		}
+
+		// ok now wait for the response, using the same buffer of the request
+		// NOTE: after using a uwsgi_buffer in that way we basically destroy (even if we can safely free it)
+		size_t rlen = ub->pos;
+        	if (uwsgi_read_with_realloc(fd, &ub->buf, &rlen, uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT])) {
+			close(fd);
+			uwsgi_buffer_destroy(ub);
+			return NULL;
+		}
+
+		// now we have a uwsgi dictionary with all of the options needed, let's parse it
+		memset(&ucmc, 0, sizeof(struct uwsgi_cache_magic_context));
+		if (uwsgi_hooked_parse(ub->buf, rlen, uwsgi_cache_magic_context_hook, &ucmc)) {
+			close(fd);
+			uwsgi_buffer_destroy(ub);
+			return NULL;
+		}
+
+		if (uwsgi_strncmp(ucmc.status, ucmc.status_len, "ok", 2)) {
+			close(fd);
+                        uwsgi_buffer_destroy(ub);
+                        return NULL;
+		}
+
+		if (ucmc.size == 0) {
+			close(fd);
+                        uwsgi_buffer_destroy(ub);
+                        return NULL;
+		}
+
+		// ok we now need to fix our buffer (if needed)
+		if (ucmc.size > rlen) {
+			char *tmp_buf = realloc(ub->buf, ucmc.size);
+			if (!tmp_buf) {
+				uwsgi_error("uwsgi_cache_magic_get()/realloc()");
+				close(fd);
+				uwsgi_buffer_destroy(ub);
+				return NULL;
+			}
+			ub->buf = tmp_buf;
+		}
+
+		if (uwsgi_read_whole_true_nb(fd, ub->buf, ucmc.size, uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT])) {
+			close(fd);
+			uwsgi_buffer_destroy(ub);
+			return NULL;
+		}
+
+		// now the magic, we dereference the internal buffer and return it to the caller
+		char *value = ub->buf;
+		ub->buf = NULL;
+		uwsgi_buffer_destroy(ub);
+		*vallen = ucmc.size;
+		return value;
+		
 	}
 
 	return NULL;

core/io.c

@@ -854,6 +854,20 @@ wait:
 }
 
 
+/*
+	like the previous one but consume the whole len (if possibile)
+*/
+
+ssize_t uwsgi_read_whole_true_nb(int fd, char *buf, size_t remains, int timeout) {
+	char *ptr = buf;
+	while(remains > 0) {
+		ssize_t len = uwsgi_read_true_nb(fd, ptr, remains, timeout);
+		if (len <= 0) return len;
+		ptr += len;
+		remains -= len;
+	}
+	return remains;
+}
 
 /*
 	this is a pretty magic function used for read a full uwsgi response
@@ -900,6 +914,7 @@ readok:
 			return -1;
 		}
 		*buffer = tmp_buf;
+		buf = *buffer;
 	}
 
 	*rlen = pktsize;

core/rpc.c

@@ -111,6 +111,7 @@ char *uwsgi_do_rpc(char *node, char *func, uint8_t argc, char *argv[], uint16_t
 	// wait for connection;
 	int ret = uwsgi.wait_write_hook(fd, uwsgi.shared->options[UWSGI_OPTION_SOCKET_TIMEOUT]);
 	if (ret <= 0) {
+		close(fd);
 		return NULL;
 	}
 

plugins/cache/cache.c

@@ -61,6 +61,46 @@ static void cache_simple_command(char *key, uint16_t keylen, char *val, uint16_t
         }
 }
 
+// this function does not use the magic api internally to avoid too much copy
+static void manage_magic_context(struct wsgi_request *wsgi_req, struct uwsgi_cache_magic_context *ucmc) {
+
+	struct uwsgi_buffer *ub = NULL;
+	struct uwsgi_cache *uc = uwsgi.caches;
+
+	if (ucmc->cache_len > 0) {
+		uc = uwsgi_cache_by_namelen(ucmc->cache, ucmc->cache_len);
+		if (!uc) return;
+	}
+
+	if (!uc) return;
+
+	// cache get
+	if (!uwsgi_strncmp(ucmc->cmd, ucmc->cmd_len, "get", 3)) {
+		uint64_t vallen = 0;
+		uwsgi_rlock(uc->lock);
+		char *value = uwsgi_cache_get2(uc, ucmc->key, ucmc->key_len, &vallen);
+		if (!value) {
+			uwsgi_rwunlock(uc->lock);
+			return;
+		}
+		// we are still locked !!!
+		ub = uwsgi_buffer_new(uwsgi.page_size);
+		ub->pos = 4;
+		if (uwsgi_buffer_append_keyval(ub, "status", 6, "ok", 2)) goto error;
+		if (uwsgi_buffer_append_keynum(ub, "size", 4, vallen)) goto error;
+		if (uwsgi_buffer_set_uh(ub, 111, 17)) goto error;
+		if (uwsgi_buffer_append(ub, value, vallen)) goto error;
+		// unlock !!!
+		uwsgi_rwunlock(uc->lock);
+		uwsgi_response_write_body_do(wsgi_req, ub->buf, ub->pos);
+		uwsgi_buffer_destroy(ub);
+		return;	
+	}
+
+error:
+	uwsgi_rwunlock(uc->lock);
+	uwsgi_buffer_destroy(ub);
+}
 
 static int uwsgi_cache_request(struct wsgi_request *wsgi_req) {
 
@@ -70,6 +110,9 @@ static int uwsgi_cache_request(struct wsgi_request *wsgi_req) {
         uint16_t argvs[3];
         uint8_t argc = 0;
 
+	// used for modifier2 17
+	struct uwsgi_cache_magic_context ucmc;
+
         switch(wsgi_req->uh->modifier2) {
                 case 0:
                         // get
@@ -156,6 +199,16 @@ static int uwsgi_cache_request(struct wsgi_request *wsgi_req) {
 			uwsgi_response_write_body_do(wsgi_req, buf, uwsgi.caches->filesize);
 			free(buf);
 			break;
+		case 17:
+			if (wsgi_req->uh->pktsize == 0) break;
+			memset(&ucmc, 0, sizeof(struct uwsgi_cache_magic_context));
+			if (uwsgi_hooked_parse(wsgi_req->buffer, wsgi_req->uh->pktsize, uwsgi_cache_magic_context_hook, &ucmc)) {
+				break;
+			}
+			manage_magic_context(wsgi_req, &ucmc);
+			break;
+		default:
+			break;
         }
 
         return UWSGI_OK;

plugins/python/python_plugin.c

@@ -832,9 +832,7 @@ void init_uwsgi_embedded_module() {
 		init_uwsgi_module_sharedarea(new_uwsgi_module);
 	}
 
-	if (uwsgi.caches) {
-		init_uwsgi_module_cache(new_uwsgi_module);
-	}
+	init_uwsgi_module_cache(new_uwsgi_module);
 
 	if (uwsgi.queue_size > 0) {
 		init_uwsgi_module_queue(new_uwsgi_module);

uwsgi.h

@@ -2681,6 +2681,7 @@ char *uwsgi_cache_get2(struct uwsgi_cache *, char *, uint16_t, uint64_t *);
 uint32_t uwsgi_cache_exists2(struct uwsgi_cache *, char *, uint16_t);
 struct uwsgi_cache *uwsgi_cache_create(char *);
 struct uwsgi_cache *uwsgi_cache_by_name(char *);
+struct uwsgi_cache *uwsgi_cache_by_namelen(char *, uint16_t);
 void uwsgi_cache_create_all(void);
 
 #define uwsgi_cache_set(x1, x2, x3, x4, x5, x6) uwsgi_cache_set2(uwsgi.caches, x1, x2, x3, x4, x5, x6)
@@ -2697,9 +2698,9 @@ void uwsgi_cache_start_sync_servers(void);
 	void *uwsgi_calloc(size_t);
 
 
-	int event_queue_init(void);
-	void *event_queue_alloc(int);
-	int event_queue_add_fd_read(int, int);
+int event_queue_init(void);
+void *event_queue_alloc(int);
+int event_queue_add_fd_read(int, int);
 	int event_queue_add_fd_write(int, int);
 	int event_queue_del_fd(int, int, int);
 	int event_queue_wait(int, int, int *);
@@ -2968,6 +2969,7 @@ void uwsgi_socket_b(int);
 int uwsgi_write_nb(int, char *, size_t, int);
 int uwsgi_read_nb(int, char *, size_t, int);
 ssize_t uwsgi_read_true_nb(int, char *, size_t, int);
+ssize_t uwsgi_read_whole_true_nb(int, char *, size_t, int);
 int uwsgi_read_uh(int fd, struct uwsgi_header *, int);
 int uwsgi_proxy_nb(struct wsgi_request *, char *, struct uwsgi_buffer *, size_t, int);
 
@@ -3596,6 +3598,7 @@ int uwsgi_buffer_insert(struct uwsgi_buffer *, size_t, char *, size_t);
 int uwsgi_buffer_insert_chunked(struct uwsgi_buffer *, size_t, size_t);
 int uwsgi_buffer_append_chunked(struct uwsgi_buffer *, size_t);
 int uwsgi_buffer_append_json(struct uwsgi_buffer *, char *, size_t);
+int uwsgi_buffer_set_uh(struct uwsgi_buffer *, uint8_t, uint8_t);
 
 ssize_t uwsgi_buffer_write_simple(struct wsgi_request *, struct uwsgi_buffer *);
 
@@ -3887,9 +3890,24 @@ struct uwsgi_stats_pusher_instance *uwsgi_stats_pusher_add(struct uwsgi_stats_pu
 
 int plugin_already_loaded(const char *);
 
+struct uwsgi_cache_magic_context {
+	char *cmd;
+	uint16_t cmd_len;
+	char *key;
+	uint16_t key_len;
+	uint64_t size;
+	uint64_t expires;
+	char *status;
+	uint16_t status_len;
+	char *cache;
+	uint16_t cache_len;
+};
+
 char *uwsgi_cache_magic_get(char *, uint16_t, uint64_t *, char *);
 int uwsgi_cache_magic_set(char *, uint16_t, char *, uint64_t, uint64_t, uint64_t, char *);
 int uwsgi_cache_magic_del(char *, uint16_t, char *);
+int uwsgi_cache_magic_exists(char *, uint16_t, char *);
+void uwsgi_cache_magic_context_hook(char *, uint16_t, char *, uint16_t, void *);
 
 #ifdef UWSGI_ZLIB
 #include <zlib.h>