diff --git a/core/buffer.c b/core/buffer.c index b9d4a087..7a244f38 100644 --- a/core/buffer.c +++ b/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; +} diff --git a/core/cache.c b/core/cache.c index 1d0f6d83..1498761b 100644 --- a/core/cache.c +++ b/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; diff --git a/core/io.c b/core/io.c index a9b69b10..1dc70637 100644 --- a/core/io.c +++ b/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; diff --git a/core/rpc.c b/core/rpc.c index 36e76d18..e4033224 100644 --- a/core/rpc.c +++ b/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; } diff --git a/plugins/cache/cache.c b/plugins/cache/cache.c index 4a4073c0..d921d488 100644 --- a/plugins/cache/cache.c +++ b/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; diff --git a/plugins/python/python_plugin.c b/plugins/python/python_plugin.c index e7c2d863..c1d0375e 100644 --- a/plugins/python/python_plugin.c +++ b/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); diff --git a/uwsgi.h b/uwsgi.h index 4c965f88..37b2858c 100644 --- a/uwsgi.h +++ b/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