proxy_ring_hash.c

/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
* Copyright (c) 2022, Cache Forge LLC, All rights reserved.
* Alan Kasindorf <alan@cacheforge.com>
* Copyright (c) 2007, Last.fm, All rights reserved.
* Richard Jones <rj@last.fm>
* Christian Muehlhaeuser <muesli@gmail.com>
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of the Last.fm Limited nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY Last.fm ``AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL Last.fm BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "proxy.h"
#include "md5.h"

#define DEFAULT_BUCKET_SIZE 160

typedef struct {
    unsigned int point; // continuum point.
    unsigned int id; // server id.
} cpoint;

typedef struct {
    struct proxy_hash_caller phc; // passed back to the proxy API.
    unsigned int total_buckets;
    cpoint continuum[]; // points to server ids.
} ketama_t;

static uint64_t ketama_key_hasher(const void *key, size_t len, uint64_t seed);
static struct proxy_hash_func ketama_key_hash = {
    ketama_key_hasher,
};

/* FROM ketama.c */
static void ketama_md5_digest( char* inString, unsigned char md5pword[16] )
{
    md5_state_t md5state;

    md5_init( &md5state );
    md5_append( &md5state, (unsigned char *)inString, strlen( inString ) );
    md5_finish( &md5state, md5pword );
}

static int ketama_compare(const void *p1, const void *p2) {
    const cpoint *a = p1;
    const cpoint *b = p2;

    return (a->point < b->point) ? -1 : ((a->point > b->point) ? 1 : 0);
}

static uint64_t ketama_key_hasher(const void *key, size_t len, uint64_t seed) {
    // NOTE: seed is ignored!
    // embedding the md5 bits since key is specified with a length here.
    md5_state_t md5state;
    unsigned char digest[16];

    md5_init(&md5state);
    md5_append(&md5state, (unsigned char *)key, len);
    md5_finish(&md5state, digest);

    // mix the hash down (from ketama_hashi)
    unsigned int h = (unsigned int)(( digest[3] << 24 )
                        | ( digest[2] << 16 )
                        | ( digest[1] <<  8 )
                        |   digest[0] );
    return h;
}

// Note: must return lookupas as zero-indexed.
static uint32_t ketama_get_server(uint64_t hash, void *ctx) {
    ketama_t *kt = (ketama_t *)ctx;
    unsigned int h = hash;
    int highp = kt->total_buckets;
    int lowp = 0, midp;
    unsigned int midval, midval1;

    // divide and conquer array search to find server with next biggest
    // point after what this key hashes to
    while ( 1 )
    {
        midp = (int)( ( lowp+highp ) / 2 );

        if ( midp == kt->total_buckets )
            return kt->continuum[0].id-1; // if at the end, roll back to zeroth

        midval = kt->continuum[midp].point;
        midval1 = midp == 0 ? 0 : kt->continuum[midp-1].point;

        if ( h <= midval && h > midval1 )
            return kt->continuum[midp].id-1;

        if ( midval < h )
            lowp = midp + 1;
        else
            highp = midp - 1;

        if ( lowp > highp )
            return kt->continuum[0].id-1;
    }
}
/* END FROM ketama.c */

// not much to be done about this without making the interface unusable.
#define MODE_DEFAULT 0 // uses xxhash
#define MODE_KETAMA 1 // uses md5
#define MODE_TWEMPROXY 2 // technically "libmemcached" ?
#define MODE_EVCACHE 3 // not sure why this has a funny string init.

// Not sure the hash algo used here matters all that much given the low number
// of points... but it might be better to let it be overrideable.
static void _add_server_default(ketama_t *kt, size_t hashstring_size, const char **parts,
        lua_Integer bucket_size, lua_Integer id, unsigned int *cont) {
    char *hashstring = malloc(hashstring_size);

    for (int k = 0; k < bucket_size; k++) {
        size_t len = snprintf(hashstring, hashstring_size, "%s:%s-%d", parts[0], parts[1], k);
        kt->continuum[*cont].point = (unsigned int) XXH3_64bits(hashstring, len);
        kt->continuum[*cont].id = id;
    }

    free(hashstring);
}

static void _add_server_ketama(ketama_t *kt, size_t hashstring_size, const char **parts,
        lua_Integer bucket_size, lua_Integer id, unsigned int *cont) {
    char *hashstring = malloc(hashstring_size);

    for (int k = 0; k < bucket_size / 4; k++) {
        unsigned char digest[16];

        // - create hashing string for ketama
        snprintf(hashstring, hashstring_size, "%s:%s-%d", parts[0], parts[1], k);

        // - md5() hash it
        // mostly from ketama.c
        ketama_md5_digest(hashstring, digest);

        /* Use successive 4-bytes from hash as numbers
         * for the points on the circle: */
        for(int h = 0; h < 4; h++ )
        {
            kt->continuum[*cont].point = ( digest[3+h*4] << 24 )
                                  | ( digest[2+h*4] << 16 )
                                  | ( digest[1+h*4] <<  8 )
                                  |   digest[h*4];
            kt->continuum[*cont].id = id;
            (*cont)++;
        }

    }

    free(hashstring);
}

static void _add_server_twemproxy(ketama_t *kt, size_t hashstring_size, const char **parts,
        lua_Integer bucket_size, lua_Integer id, unsigned int *cont) {
    char *hashstring = malloc(hashstring_size);

    for (int k = 0; k < bucket_size / 4; k++) {
        unsigned char digest[16];

        // - create hashing string for ketama
        if (strcmp(parts[1], "11211") == 0) {
            // twemproxy sources libmemcached as removing the default port
            // from the string if found.
            snprintf(hashstring, hashstring_size, "%s-%d", parts[0], k);
        } else {
            snprintf(hashstring, hashstring_size, "%s:%s-%d", parts[0], parts[1], k);
        }

        // - md5() hash it
        // mostly from ketama.c
        ketama_md5_digest(hashstring, digest);

        /* Use successive 4-bytes from hash as numbers
         * for the points on the circle: */
        for(int h = 0; h < 4; h++ )
        {
            kt->continuum[*cont].point = ( digest[3+h*4] << 24 )
                                  | ( digest[2+h*4] << 16 )
                                  | ( digest[1+h*4] <<  8 )
                                  |   digest[h*4];
            kt->continuum[*cont].id = id;
            (*cont)++;
        }

    }

    free(hashstring);
}

static void _add_server_evcache(ketama_t *kt, size_t hashstring_size, const char **parts,
        lua_Integer bucket_size, lua_Integer id, unsigned int *cont) {
    char *hashstring = malloc(hashstring_size);

    for (int k = 0; k < bucket_size / 4; k++) {
        unsigned char digest[16];

        // - create hashing string for ketama
        snprintf(hashstring, hashstring_size, "%s/%s:%s-%d", parts[0], parts[0], parts[1], k);
        // - md5() hash it
        // mostly from ketama.c
        ketama_md5_digest(hashstring, digest);

        /* Use successive 4-bytes from hash as numbers
         * for the points on the circle: */
        for(int h = 0; h < 4; h++ )
        {
            kt->continuum[*cont].point = ( digest[3+h*4] << 24 )
                                  | ( digest[2+h*4] << 16 )
                                  | ( digest[1+h*4] <<  8 )
                                  |   digest[h*4];
            kt->continuum[*cont].id = id;
            (*cont)++;
        }

    }

    free(hashstring);
}

#define PARTS 2
// stack = [pool, option]
static int ketama_new(lua_State *L) {
    lua_Integer bucket_size = DEFAULT_BUCKET_SIZE;
    const char *parts[PARTS];
    size_t partlens[PARTS];
    int makemode = 0;

    // check for UA_TTABLE at 1
    luaL_checktype(L, 1, LUA_TTABLE);
    // get number of servers in pool.
    // NOTE: rawlen skips metatable redirection. if we care; lua_len instead.
    lua_Unsigned total = lua_rawlen(L, 1);

    // check for optional input
    int argc = lua_gettop(L);
    if (argc > 1) {
        luaL_checktype(L, 2, LUA_TTABLE);
        if (lua_getfield(L, 2, "omode") != LUA_TNIL) {
            luaL_checktype(L, -1, LUA_TSTRING);
            const char *mode = lua_tostring(L, -1);
            if (strcmp(mode, "default") == 0) {
                makemode = MODE_DEFAULT;
            } else if (strcmp(mode, "ketama") == 0) {
                makemode = MODE_KETAMA;
            } else if (strcmp(mode, "twemproxy") == 0) {
                makemode = MODE_TWEMPROXY;
            } else if (strcmp(mode, "evcache") == 0) {
                makemode = MODE_EVCACHE;
            } else {
                lua_pushstring(L, "ring_hash: bad omode argument");
                lua_error(L);
            }
        }
        lua_pop(L, 1); // pops the nil or mode

        if (lua_getfield(L, 2, "obuckets") != LUA_TNIL) {
          int success = 0;
          bucket_size = lua_tointegerx(L, -1, &success);
          if (!success || bucket_size < 1) {
              lua_pushstring(L, "ring_hash: option argument must be a positive integer");
              lua_error(L);
          }
        }
        lua_pop(L, 1);
    }

    // newuserdatauv() sized for pool*
    size_t size = sizeof(ketama_t) + sizeof(cpoint) * (total * bucket_size);
    ketama_t *kt = lua_newuserdatauv(L, size, 0);
    // TODO: check *kt.
    kt->total_buckets = bucket_size * total;

    // loop over pool
    unsigned int cont = 0;
    lua_pushnil(L); // start the pool iterator
    while (lua_next(L, 1) != 0) {
        // key is -2, value is -1.
        // value is another table. need to query it to get what we need for
        // the hash.
        // hash string is: hostname/ipaddr:port-repitition
        // TODO: bother doing error checking?
        lua_getfield(L, -1, "id");
        lua_Integer id = lua_tointeger(L, -1);
        lua_pop(L, 1);

        // FIXME: we need to do the lua_pop after string assembly to be safe.
        lua_getfield(L, -1, "addr");
        parts[0] = lua_tolstring(L, -1, &partlens[0]);
        lua_pop(L, 1);
        lua_getfield(L, -1, "port");
        parts[1] = lua_tolstring(L, -1, &partlens[1]);
        lua_pop(L, 1);

        size_t hashstring_size = 0;
        for (int x = 0; x < PARTS; x++) {
            hashstring_size += partlens[x];
        }

        // We have up to 3 delimiters in the final hashstring and an index
        // 16 bytes is plenty to accomodate this requirement.
        hashstring_size += 16;

        switch (makemode) {
            case MODE_DEFAULT:
                _add_server_default(kt, hashstring_size, parts, bucket_size, id, &cont);
                break;
            case MODE_KETAMA:
                _add_server_ketama(kt, hashstring_size, parts, bucket_size, id, &cont);
                break;
            case MODE_TWEMPROXY:
                _add_server_twemproxy(kt, hashstring_size, parts, bucket_size, id, &cont);
                break;
            case MODE_EVCACHE:
                _add_server_evcache(kt, hashstring_size, parts, bucket_size, id, &cont);
                break;
        }

        lua_pop(L, 1); // remove value, leave key for next iteration.
    }

    // - qsort the points
    qsort( &kt->continuum, cont, sizeof(cpoint), ketama_compare);

    // set the hash/fetch function and the context ptr.
    kt->phc.ctx = kt;
    kt->phc.selector_func = ketama_get_server;

    // - add a pushlightuserdata for the sub-struct with func/ctx.
    lua_pushlightuserdata(L, &kt->phc);
    // - return [UD, lightuserdata]
    return 2;
}

int mcplib_open_dist_ring_hash(lua_State *L) {
    const struct luaL_Reg ketama_f[] = {
        {"new", ketama_new},
        {NULL, NULL},
    };

    luaL_newlib(L, ketama_f);
    lua_pushlightuserdata(L, &ketama_key_hash);
    lua_setfield(L, -2, "hash");

    return 1;
}