teems/Enclave/route.cpp

#include "Enclave_t.h"
#include "config.hpp"
#include "utils.hpp"
#include "sort.hpp"
#include "comms.hpp"
#include "obliv.hpp"
#include "storage.hpp"
#include "route.hpp"

#define PROFILE_ROUTING
#define TRACE_ROUTING

RouteState route_state;

// Computes ceil(x/y) where x and y are integers, x>=0, y>0.

#define CEILDIV(x,y) (((x)+(y)-1)/(y))

#ifdef TRACE_ROUTING

// Show (the first 300 of, if there are more) the headers and the first
// few bytes of the body of each message in the buffer
static void show_messages(const char *label, const unsigned char *buffer,
    size_t num)
{
    if (label) {
        printf("%s\n", label);
    }
    for (size_t i=0; i<num && i<300; ++i) {
        const uint32_t *ibuf = (const uint32_t *)buffer;
        if (g_teems_config.private_routing) {
            printf("%3d R:%08x S:%08x [%08x]\n", i, ibuf[0], ibuf[1],
                ibuf[2]);
        } else {
            printf("%3d R:%08x P:%08x S:%08x [%08x]\n", i, ibuf[0], ibuf[1],
                ibuf[2], ibuf[3]);
        }
        buffer += g_teems_config.msg_size;
    }
    printf("\n");
}

#endif

// Call this near the end of ecall_config_load, but before
// comms_init_nodestate. Returns true on success, false on failure.
bool route_init()
{
    // Compute the maximum number of messages we could receive by direct
    // ingestion

    // Each ingestion node will have at most
    // ceil(user_count/num_ingestion_nodes) users, and each user will
    // send at most m_priv_out messages.
    uint32_t users_per_ing = CEILDIV(g_teems_config.user_count,
        g_teems_config.num_ingestion_nodes);
    uint32_t tot_msg_per_ing;
    if (g_teems_config.private_routing) {
        tot_msg_per_ing = users_per_ing * g_teems_config.m_priv_out;
    } else {
        tot_msg_per_ing = users_per_ing * g_teems_config.m_pub_out;
    }

    // Compute the maximum number of messages we could receive in round 1
    // In private routing, each ingestion node will send us an
    // our_weight/tot_weight fraction of the messages they hold
    uint32_t max_msg_from_each_ing;
    max_msg_from_each_ing = CEILDIV(tot_msg_per_ing, g_teems_config.tot_weight) *
        g_teems_config.my_weight;

    // And the maximum number we can receive in total is that times the
    // number of ingestion nodes
    uint32_t max_round1_msgs = max_msg_from_each_ing *
        g_teems_config.num_ingestion_nodes;

    // Compute the maximum number of messages we could send in round 2

    // Each storage node has at most this many users
    uint32_t users_per_stg = CEILDIV(g_teems_config.user_count,
        g_teems_config.num_storage_nodes);

    // And so can receive at most this many messages
    uint32_t tot_msg_per_stg;
    if (g_teems_config.private_routing) {
        tot_msg_per_stg = users_per_stg * g_teems_config.m_priv_in;
    } else {
        tot_msg_per_stg = users_per_stg * g_teems_config.m_pub_in;
    }

    // Which will be at most this many from us
    uint32_t max_msg_to_each_stg;
    max_msg_to_each_stg = CEILDIV(tot_msg_per_stg, g_teems_config.tot_weight) *
        g_teems_config.my_weight;

    // But we can't send more messages to each storage server than we
    // could receive in total
    if (max_msg_to_each_stg > max_round1_msgs) {
        max_msg_to_each_stg = max_round1_msgs;
    }

    // And the max total number of outgoing messages in round 2 is then
    uint32_t max_round2_msgs = max_msg_to_each_stg *
        g_teems_config.num_storage_nodes;

    // In case we have a weird configuration where users can send more
    // messages per epoch than they can receive, ensure the round 2
    // buffer is large enough to hold the incoming messages as well
    if (max_round2_msgs < max_round1_msgs) {
        max_round2_msgs = max_round1_msgs;
    }

    // The max number of messages that can arrive at a storage server
    uint32_t max_stg_msgs;
    max_stg_msgs = tot_msg_per_stg + g_teems_config.tot_weight;

    // Calculating public-routing buffer sizes
    // Weights are not used in public routing
    uint32_t max_round1b_msgs_to_adj_rtr =
        (g_teems_config.num_routing_nodes-1)*(g_teems_config.num_routing_nodes-1);
    // Ensure columnroute constraint that column height is >= 2*(num_routing_nodes-1)^2
    // and round up to a multiple of num_routing_nodes
    uint32_t max_round1a_msgs = CEILDIV(
        std::max(max_round1_msgs, 2*max_round1b_msgs_to_adj_rtr),
        g_teems_config.num_routing_nodes) *
        g_teems_config.num_routing_nodes;
    uint32_t max_round1c_msgs = max_round1a_msgs;

    /*
    printf("users_per_ing=%u, tot_msg_per_ing=%u, max_msg_from_each_ing=%u, max_round1_msgs=%u, users_per_stg=%u, tot_msg_per_stg=%u, max_msg_to_each_stg=%u, max_round2_msgs=%u, max_stg_msgs=%u\n", users_per_ing, tot_msg_per_ing, max_msg_from_each_ing, max_round1_msgs, users_per_stg, tot_msg_per_stg, max_msg_to_each_stg, max_round2_msgs, max_stg_msgs);
    */

#ifdef TRACK_HEAP_USAGE
    printf("route_init H1 heap %u\n", g_peak_heap_used);
#endif
    // Create the route state
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];
    try {
        if (my_roles & ROLE_INGESTION) {
            route_state.ingbuf.alloc(tot_msg_per_ing);
        }
#ifdef TRACK_HEAP_USAGE
        printf("route_init alloc %u msgs\n", tot_msg_per_ing);
        printf("route_init H2 heap %u\n", g_peak_heap_used);
#endif
        if (my_roles & ROLE_ROUTING) {
            route_state.round1.alloc(max_round2_msgs);
#ifdef TRACK_HEAP_USAGE
            printf("route_init alloc %u msgs\n", max_round2_msgs);
            printf("route_init H3 heap %u\n", g_peak_heap_used);
#endif
            if (!g_teems_config.private_routing) {
                route_state.round1a.alloc(max_round1a_msgs);
                route_state.round1a_sorted.alloc(max_round1a_msgs);
                // double round 1b buffers to sort with some round 1a messages
                route_state.round1b_prev.alloc(2*max_round1b_msgs_to_adj_rtr);
                route_state.round1b_next.alloc(2*max_round1b_msgs_to_adj_rtr);
                route_state.round1c.alloc(max_round1c_msgs);
#ifdef TRACK_HEAP_USAGE
                printf("route_init alloc %u msgs\n", max_round1a_msgs);
                printf("route_init alloc %u msgs\n", max_round1a_msgs);
                printf("route_init alloc %u msgs\n", 2*max_round1b_msgs_to_adj_rtr);
                printf("route_init alloc %u msgs\n", 2*max_round1b_msgs_to_adj_rtr);
                printf("route_init alloc %u msgs\n", max_round1c_msgs);
                printf("route_init H4 heap %u\n", g_peak_heap_used);
#endif
            }
        }
        if (my_roles & ROLE_STORAGE) {
            route_state.round2.alloc(max_stg_msgs);
#ifdef TRACK_HEAP_USAGE
                printf("route_init alloc %u msgs\n", max_stg_msgs);
                printf("route_init H5 heap %u\n", g_peak_heap_used);
#endif
            if (!storage_init(users_per_stg, max_stg_msgs)) {
                return false;
            }
        }
    } catch (std::bad_alloc&) {
        printf("Memory allocation failed in route_init\n");
        return false;
    }
    route_state.step = ROUTE_NOT_STARTED;
    route_state.tot_msg_per_ing = tot_msg_per_ing;
    route_state.max_round1_msgs = max_round1_msgs;
    route_state.max_round1a_msgs = max_round1a_msgs;
    route_state.max_round1b_msgs_to_adj_rtr = max_round1b_msgs_to_adj_rtr;
    route_state.max_round1c_msgs = max_round1c_msgs;
    route_state.max_msg_to_each_stg = max_msg_to_each_stg;
    route_state.max_round2_msgs = max_round2_msgs;
    route_state.max_stg_msgs = max_stg_msgs;
    route_state.cbpointer = NULL;

    threadid_t nthreads = g_teems_config.nthreads;
#ifdef PROFILE_ROUTING
    unsigned long start = printf_with_rtclock("begin precompute evalplans (%u,%hu) (%u,%hu)\n", tot_msg_per_ing, nthreads, max_round2_msgs, nthreads);
#endif
    if (my_roles & ROLE_INGESTION) {
        sort_precompute_evalplan(tot_msg_per_ing, nthreads);
    }
    if (my_roles & ROLE_ROUTING) {
        sort_precompute_evalplan(max_round2_msgs, nthreads);
        if(!g_teems_config.private_routing) {
            sort_precompute_evalplan(max_round1a_msgs, nthreads);
            sort_precompute_evalplan(2*max_round1b_msgs_to_adj_rtr, nthreads);
        }
    }
    if (my_roles & ROLE_STORAGE) {
        sort_precompute_evalplan(max_stg_msgs, nthreads);
    }
#ifdef PROFILE_ROUTING
    printf_with_rtclock_diff(start, "end precompute evalplans\n");
#endif
#ifdef TRACK_HEAP_USAGE
    printf("route_init end heap %u\n", g_peak_heap_used);
#endif
    return true;
}

// Call when shutting system down to deallocate routing state
void route_close() {
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];
    if (my_roles & ROLE_STORAGE) {
        storage_close();
    }
}

// Precompute the WaksmanNetworks needed for the sorts.  If you pass -1,
// it will return the number of different sizes it needs to regenerate.
// If you pass [0,sizes-1], it will compute one WaksmanNetwork with that
// size index and return the number of available WaksmanNetworks of that
// size.  If you pass anything else, it will return the number of
// different sizes it needs at all.

// The list of sizes that need refilling, updated when you pass -1
static std::vector<uint32_t> used_sizes;

size_t ecall_precompute_sort(int sizeidx)
{
    size_t ret = 0;

    if (sizeidx == -1) {
        used_sizes = sort_get_used();
        ret = used_sizes.size();
    } else if (sizeidx >= 0 && sizeidx < used_sizes.size()) {
        uint32_t size = used_sizes[sizeidx];
#ifdef PROFILE_ROUTING
        unsigned long start = printf_with_rtclock("begin precompute WaksmanNetwork (%u)\n", size);
#endif
        ret = sort_precompute(size);
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start, "end precompute Waksman Network (%u)\n", size);
#endif
    } else {
        uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];

        if (my_roles & ROLE_INGESTION) {
            used_sizes.push_back(route_state.tot_msg_per_ing);
        }
        if (my_roles & ROLE_ROUTING) {
            used_sizes.push_back(route_state.max_round2_msgs);
            if(!g_teems_config.private_routing) {
                used_sizes.push_back(route_state.max_round1a_msgs);
                used_sizes.push_back(2*route_state.max_round1b_msgs_to_adj_rtr);
                used_sizes.push_back(2*route_state.max_round1b_msgs_to_adj_rtr);
                used_sizes.push_back(route_state.max_round1c_msgs);
            }
        }
        if (my_roles & ROLE_STORAGE) {
            used_sizes.push_back(route_state.max_stg_msgs);
            if(!g_teems_config.private_routing) {
                used_sizes.push_back(route_state.max_stg_msgs);
            }
        }
        ret = used_sizes.size();
    }
    return ret;
}

static uint8_t* msgbuffer_get_buf(MsgBuffer &msgbuf,
        NodeCommState &, uint32_t tot_enc_chunk_size)
{
    uint16_t msg_size = g_teems_config.msg_size;

    // Chunks will be encrypted and have a MAC tag attached which will
    // not correspond to plaintext bytes, so we can trim them.

    // The minimum number of chunks needed to transmit this message
    uint32_t min_num_chunks =
        (tot_enc_chunk_size + (FRAME_SIZE-1)) / FRAME_SIZE;
    // The number of plaintext bytes this message could contain
    uint32_t plaintext_bytes = tot_enc_chunk_size -
        SGX_AESGCM_MAC_SIZE * min_num_chunks;

    assert ((plaintext_bytes % uint32_t(msg_size)) == 0);

    uint32_t num_msgs = plaintext_bytes/uint32_t(msg_size);

    pthread_mutex_lock(&msgbuf.mutex);
    uint32_t start = msgbuf.reserved;
    if (start + num_msgs > msgbuf.bufsize) {
        pthread_mutex_unlock(&msgbuf.mutex);
        printf("Max %u messages exceeded\n", msgbuf.bufsize);
        return NULL;
    }
    msgbuf.reserved += num_msgs;
    pthread_mutex_unlock(&msgbuf.mutex);

    return msgbuf.buf + start * msg_size;
}

static void round1a_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

static void round1b_prev_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

static void round1b_next_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

static void round1c_received(NodeCommState &nodest, uint8_t *data,
    uint32_t plaintext_len, uint32_t);

static void round2_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

// A round 1 message was received by a routing node from an ingestion
// node; we put it into the round 2 buffer for processing in round 2
static void round1_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[nodest.node_num];

    pthread_mutex_lock(&route_state.round1.mutex);
    route_state.round1.inserted += num_msgs;
    route_state.round1.nodes_received += 1;
    nodenum_t nodes_received = route_state.round1.nodes_received;
    bool completed_prev_round = route_state.round1.completed_prev_round;
    pthread_mutex_unlock(&route_state.round1.mutex);

    // What is the next message we expect from this node?
    if (g_teems_config.private_routing) {
        if ((our_roles & ROLE_STORAGE) && (their_roles & ROLE_ROUTING)) {
            nodest.in_msg_get_buf = [&](NodeCommState &commst,
                    uint32_t tot_enc_chunk_size) {
                return msgbuffer_get_buf(route_state.round2, commst,
                    tot_enc_chunk_size);
            };
            nodest.in_msg_received = round2_received;
        }
        // Otherwise, it's just the next round 1 message, so don't change
        // the handlers.
    } else {
        if (their_roles & ROLE_ROUTING) {
            nodest.in_msg_get_buf = [&](NodeCommState &commst,
                    uint32_t tot_enc_chunk_size) {
                return msgbuffer_get_buf(route_state.round1a, commst,
                    tot_enc_chunk_size);
            };
            nodest.in_msg_received = round1a_received;
        }
        // Otherwise, it's just the next round 1 message, so don't change
        // the handlers.
    }

    if (nodes_received == g_teems_config.num_ingestion_nodes &&
            completed_prev_round) {
        route_state.step = ROUTE_ROUND_1;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, 1);
    }
}

// A round 1a message was received by a routing node
static void round1a_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    pthread_mutex_lock(&route_state.round1a.mutex);
    route_state.round1a.inserted += num_msgs;
    route_state.round1a.nodes_received += 1;
    nodenum_t nodes_received = route_state.round1a.nodes_received;
    bool completed_prev_round = route_state.round1a.completed_prev_round;
    pthread_mutex_unlock(&route_state.round1a.mutex);

    // Both are routing nodes
    // We only expect a message from the previous and next nodes (if they exist)
    nodenum_t my_node_num = g_teems_config.my_node_num;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;
    if ((prev_nodes > 0) &&
        (nodest.node_num == g_teems_config.routing_nodes[num_routing_nodes-1])) {
        // Node is previous routing node
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1b_prev, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1b_prev_received;
    } else if ((prev_nodes < num_routing_nodes-1) &&
        (nodest.node_num == g_teems_config.routing_nodes[1])) {
        // Node is next routing node
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1b_next, commst,
                tot_enc_chunk_size);
        };        
        nodest.in_msg_received = round1b_next_received;
    } else {
        // other routing nodes will not send to this node until round 1c
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1c, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1c_received;
    }

    if (nodes_received == g_teems_config.num_routing_nodes &&
            completed_prev_round) {
        route_state.step = ROUTE_ROUND_1A;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, ROUND_1A);
    }
}

// A round 1b message was received from the previous routing node
static void round1b_prev_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[nodest.node_num];
    pthread_mutex_lock(&route_state.round1b_prev.mutex);
    route_state.round1b_prev.inserted += num_msgs;
    route_state.round1b_prev.nodes_received += 1;
    nodenum_t nodes_received = route_state.round1b_prev.nodes_received;
    bool completed_prev_round = route_state.round1b_prev.completed_prev_round;
    pthread_mutex_unlock(&route_state.round1b_prev.mutex);
    pthread_mutex_lock(&route_state.round1b_next.mutex);
    nodes_received += route_state.round1b_next.nodes_received;
    pthread_mutex_unlock(&route_state.round1b_next.mutex);
    nodest.in_msg_get_buf = [&](NodeCommState &commst,
            uint32_t tot_enc_chunk_size) {
        return msgbuffer_get_buf(route_state.round1c, commst,
            tot_enc_chunk_size);
    };
    nodest.in_msg_received = round1c_received;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    nodenum_t adjacent_nodes;
    if (num_routing_nodes == 1) {
        adjacent_nodes = 0;
    } else if ((prev_nodes == 0) || (prev_nodes == num_routing_nodes-1)) {
        adjacent_nodes = 1;
    } else {
        adjacent_nodes = 2;
    }
    if (nodes_received == adjacent_nodes && completed_prev_round) {
        route_state.step = ROUTE_ROUND_1B;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, ROUND_1B);
    }
}

// A round 1b message was received from the next routing node
static void round1b_next_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[nodest.node_num];
    pthread_mutex_lock(&route_state.round1b_next.mutex);
    route_state.round1b_next.inserted += num_msgs;
    route_state.round1b_next.nodes_received += 1;
    nodenum_t nodes_received = route_state.round1b_next.nodes_received;
    bool completed_prev_round = route_state.round1b_next.completed_prev_round;
    pthread_mutex_unlock(&route_state.round1b_next.mutex);
    pthread_mutex_lock(&route_state.round1b_prev.mutex);
    nodes_received += route_state.round1b_prev.nodes_received;
    pthread_mutex_unlock(&route_state.round1b_prev.mutex);
    nodest.in_msg_get_buf = [&](NodeCommState &commst,
            uint32_t tot_enc_chunk_size) {
        return msgbuffer_get_buf(route_state.round1c, commst,
            tot_enc_chunk_size);
    };
    nodest.in_msg_received = round1c_received;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    nodenum_t adjacent_nodes;
    if (num_routing_nodes == 1) {
        adjacent_nodes = 0;
    } else if ((prev_nodes == 0) || (prev_nodes == num_routing_nodes-1)) {
        adjacent_nodes = 1;
    } else {
        adjacent_nodes = 2;
    }
    if (nodes_received == adjacent_nodes && completed_prev_round) {
        route_state.step = ROUTE_ROUND_1B;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, ROUND_1B);
    }
}

// Message received in round 1c
static void round1c_received(NodeCommState &nodest, uint8_t *data,
    uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[nodest.node_num];
    pthread_mutex_lock(&route_state.round1c.mutex);
    route_state.round1c.inserted += num_msgs;
    route_state.round1c.nodes_received += 1;
    nodenum_t nodes_received = route_state.round1c.nodes_received;
    bool completed_prev_round = route_state.round1c.completed_prev_round;
    pthread_mutex_unlock(&route_state.round1c.mutex);

    // What is the next message we expect from this node?
    if (our_roles & ROLE_STORAGE) {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round2, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round2_received;
    } else if (their_roles & ROLE_INGESTION) {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1_received;
    } else {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1a, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1a_received;
    }
    if (nodes_received == g_teems_config.num_routing_nodes &&
            completed_prev_round) {
        route_state.step = ROUTE_ROUND_1C;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, ROUND_1C);
    }
}

// A round 2 message was received by a storage node from a routing node
static void round2_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    uint16_t msg_size = g_teems_config.msg_size;
    assert((plaintext_len % uint32_t(msg_size)) == 0);
    uint32_t num_msgs = plaintext_len / uint32_t(msg_size);
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[nodest.node_num];

    pthread_mutex_lock(&route_state.round2.mutex);
    route_state.round2.inserted += num_msgs;
    route_state.round2.nodes_received += 1;
    nodenum_t nodes_received = route_state.round2.nodes_received;
    bool completed_prev_round = route_state.round2.completed_prev_round;
    pthread_mutex_unlock(&route_state.round2.mutex);

    // What is the next message we expect from this node?
    if ((our_roles & ROLE_ROUTING) && (their_roles & ROLE_INGESTION)) {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1_received;
    }
    // Otherwise, it's just the next round 2 message, so don't change
    // the handlers.

    if (nodes_received == g_teems_config.num_routing_nodes &&
            completed_prev_round) {
        route_state.step = ROUTE_ROUND_2;
        void *cbpointer = route_state.cbpointer;
        route_state.cbpointer = NULL;
        ocall_routing_round_complete(cbpointer, 2);
    }
}

// For a given other node, set the received message handler to the first
// message we would expect from them, given their roles and our roles.
void route_init_msg_handler(nodenum_t node_num)
{
    // Our roles and their roles
    uint8_t our_roles = g_teems_config.roles[g_teems_config.my_node_num];
    uint8_t their_roles = g_teems_config.roles[node_num];

    // The node communication state
    NodeCommState &nodest = g_commstates[node_num];

    // If we are a routing node (possibly among other roles) and they
    // are an ingestion node (possibly among other roles), a round 1
    // routing message is the first thing we expect from them.  We put
    // these messages into the round1 buffer for processing.
    if ((our_roles & ROLE_ROUTING) && (their_roles & ROLE_INGESTION)) {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round1, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round1_received;
    }
    // Otherwise, if we are a storage node (possibly among other roles)
    // and they are a routing node (possibly among other roles), a round
    // 2 routing message is the first thing we expect from them
    else if ((our_roles & ROLE_STORAGE) && (their_roles & ROLE_ROUTING)) {
        nodest.in_msg_get_buf = [&](NodeCommState &commst,
                uint32_t tot_enc_chunk_size) {
            return msgbuffer_get_buf(route_state.round2, commst,
                tot_enc_chunk_size);
        };
        nodest.in_msg_received = round2_received;
    }
    // Otherwise, we don't expect a message from this node. Set the
    // unknown message handler.
    else {
        nodest.in_msg_get_buf = default_in_msg_get_buf;
        nodest.in_msg_received = unknown_in_msg_received;
    }
}

// Directly ingest a buffer of num_msgs messages into the ingbuf buffer.
// Return true on success, false on failure.
bool ecall_ingest_raw(uint8_t *msgs, uint32_t num_msgs)
{
    uint16_t msg_size = g_teems_config.msg_size;
    MsgBuffer &ingbuf = route_state.ingbuf;

    pthread_mutex_lock(&ingbuf.mutex);
    uint32_t start = ingbuf.reserved;
    if (start + num_msgs > route_state.tot_msg_per_ing) {
        pthread_mutex_unlock(&ingbuf.mutex);
        printf("Max %u messages exceeded\n",
            route_state.tot_msg_per_ing);
        return false;
    }
    ingbuf.reserved += num_msgs;
    pthread_mutex_unlock(&ingbuf.mutex);

    memmove(ingbuf.buf + start * msg_size,
        msgs, num_msgs * msg_size);

    pthread_mutex_lock(&ingbuf.mutex);
    ingbuf.inserted += num_msgs;
    pthread_mutex_unlock(&ingbuf.mutex);

    return true;
}

// Send messages round-robin, used in rounds 1 and 1c.  Note that N here is not private.
template<typename T>
static void send_round_robin_msgs(MsgBuffer &round, const uint8_t *msgs, const T *indices,
    uint32_t N)
{
    uint16_t msg_size = g_teems_config.msg_size;
    uint16_t tot_weight;
    tot_weight = g_teems_config.tot_weight;
    nodenum_t my_node_num = g_teems_config.my_node_num;

    uint32_t full_rows;
    uint32_t last_row;
    full_rows = N / uint32_t(tot_weight);
    last_row = N % uint32_t(tot_weight);

    for (auto &routing_node: g_teems_config.routing_nodes) {
        uint8_t weight = g_teems_config.weights[routing_node].weight;

        if (weight == 0) {
            // This shouldn't happen, but just in case
            continue;
        }
        uint16_t start_weight =
            g_teems_config.weights[routing_node].startweight;

        // The number of messages headed for this routing node from the
        // full rows
        uint32_t num_msgs_full_rows = full_rows * uint32_t(weight);
        // The number of messages headed for this routing node from the
        // incomplete last row is:
        // 0 if last_row < start_weight
        // last_row-start_weight if start_weight <= last_row < start_weight + weight
        // weight if start_weight + weight <= last_row
        uint32_t num_msgs_last_row = 0;
        if (start_weight <= last_row && last_row < start_weight + weight) {
            num_msgs_last_row = last_row-start_weight;
        } else if (start_weight + weight <= last_row) {
            num_msgs_last_row = weight;
        }
        // The total number of messages headed for this routing node
        uint32_t num_msgs = num_msgs_full_rows + num_msgs_last_row;

        if (routing_node == my_node_num) {
            // Special case: we're sending to ourselves; just put the
            // messages in our own buffer

            pthread_mutex_lock(&round.mutex);
            uint32_t start = round.reserved;
            if (start + num_msgs > round.bufsize) {
                pthread_mutex_unlock(&round.mutex);
                printf("Max %u messages exceeded\n", round.bufsize);
                return;
            }
            round.reserved += num_msgs;
            pthread_mutex_unlock(&round.mutex);
            uint8_t *buf = round.buf + start * msg_size;

            for (uint32_t i=0; i<full_rows; ++i) {
                const T *idxp = indices + i*tot_weight + start_weight;
                for (uint32_t j=0; j<weight; ++j) {
                    memmove(buf, msgs + idxp[j].index()*msg_size, msg_size);
                    buf += msg_size;
                }
            }
            const T *idxp = indices + full_rows*tot_weight + start_weight;
            for (uint32_t j=0; j<num_msgs_last_row; ++j) {
                memmove(buf, msgs + idxp[j].index()*msg_size, msg_size);
                buf += msg_size;
            }

            pthread_mutex_lock(&round.mutex);
            round.inserted += num_msgs;
            round.nodes_received += 1;
            pthread_mutex_unlock(&round.mutex);

        } else {
            NodeCommState &nodecom = g_commstates[routing_node];
            nodecom.message_start(num_msgs * msg_size);
            for (uint32_t i=0; i<full_rows; ++i) {
                const T *idxp = indices + i*tot_weight + start_weight;
                for (uint32_t j=0; j<weight; ++j) {
                    nodecom.message_data(msgs + idxp[j].index()*msg_size, msg_size);
                }
            }
            const T *idxp = indices + full_rows*tot_weight + start_weight;
            for (uint32_t j=0; j<num_msgs_last_row; ++j) {
                nodecom.message_data(msgs + idxp[j].index()*msg_size, msg_size);
            }
        }
    }
}

// Send the round 1a messages from the round 1 buffer, which only occurs in public-channel routing.
// msgs points to the message buffer, indices points to the the sorted indices, and N is the number
// of non-padding items.
static void send_round1a_msgs(const uint8_t *msgs, const UidPriorityKey *indices, uint32_t N) {
    uint16_t msg_size = g_teems_config.msg_size;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;

    uint32_t min_msgs_per_node = route_state.max_round1a_msgs / num_routing_nodes;
    uint32_t extra_msgs = route_state.max_round1a_msgs % num_routing_nodes;
    for (auto &routing_node: g_teems_config.routing_nodes) {
        // In this unweighted setting, start_weight represents the position among routing nodes
        uint16_t prev_nodes = g_teems_config.weights[routing_node].startweight;
        uint32_t start_msg, num_msgs;
        if (prev_nodes >= extra_msgs) {
            start_msg = min_msgs_per_node * prev_nodes + extra_msgs;
            num_msgs = min_msgs_per_node;
        } else {
            start_msg = min_msgs_per_node * prev_nodes + prev_nodes;
            num_msgs = min_msgs_per_node + 1;
        }
        // take number of messages into account
        if (start_msg >= N) {
            num_msgs = 0;
        } else if (start_msg + num_msgs > N) {
            num_msgs = N - start_msg;
        }

        if (routing_node == my_node_num) {
            // Special case: we're sending to ourselves; just put the
            // messages in our own buffer
            MsgBuffer &round1a = route_state.round1a;
            pthread_mutex_lock(&round1a.mutex);
            uint32_t start = round1a.reserved;
            if (start + num_msgs > round1a.bufsize) {
                pthread_mutex_unlock(&round1a.mutex);
                printf("Max %u messages exceeded in round 1a\n", round1a.bufsize);
                return;
            }
            round1a.reserved += num_msgs;
            pthread_mutex_unlock(&round1a.mutex);
            uint8_t *buf = round1a.buf + start * msg_size;
            for (uint32_t i=0; i<num_msgs; ++i) {
                const UidPriorityKey *idxp = indices + start_msg + i;
                memmove(buf, msgs + idxp->index()*msg_size, msg_size);
                buf += msg_size;
            }
            pthread_mutex_lock(&round1a.mutex);
            round1a.inserted += num_msgs;
            round1a.nodes_received += 1;
            pthread_mutex_unlock(&round1a.mutex);
        } else {
            NodeCommState &nodecom = g_commstates[routing_node];
            nodecom.message_start(num_msgs * msg_size);
            for (uint32_t i=0; i<num_msgs; ++i) {
                const UidPriorityKey *idxp = indices + start_msg + i;
                nodecom.message_data(msgs + idxp->index()*msg_size, msg_size);
            }
        }
    }
}

// Send the round 1b messages from the round 1a buffer, which only occurs in public-channel routing.
// msgs points to the message buffer, and N is the number of non-padding items.
static void send_round1b_msgs(const uint8_t *msgs, uint32_t N) {
    uint16_t msg_size = g_teems_config.msg_size;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;

    // send to previous node
    if (prev_nodes > 0) {
        nodenum_t prev_node = g_teems_config.routing_nodes[num_routing_nodes-1];
        NodeCommState &nodecom = g_commstates[prev_node];
        uint32_t num_msgs = min(N, route_state.max_round1b_msgs_to_adj_rtr);
        nodecom.message_start(num_msgs * msg_size);
        for (uint32_t i=0; i<num_msgs; ++i) {
            nodecom.message_data(msgs + i*msg_size, msg_size);
        }
    }
    // send to next node
    if (prev_nodes < num_routing_nodes-1) {
        nodenum_t next_node = g_teems_config.routing_nodes[1];
        NodeCommState &nodecom = g_commstates[next_node];
        if (N <= route_state.max_round1a_msgs - route_state.max_round1b_msgs_to_adj_rtr) {
            // No messages to exchange with next node
            nodecom.message_start(0);
            // No need to call message_data()
        } else {
            uint32_t start_msg =
                route_state.max_round1a_msgs - route_state.max_round1b_msgs_to_adj_rtr;
            uint32_t num_msgs = N - start_msg;
            nodecom.message_start(num_msgs * msg_size);
            for (uint32_t i=0; i<num_msgs; ++i) {
                nodecom.message_data(msgs + i*msg_size, msg_size);
            }
        }
    }
}

// Send the round 2 messages from the previous-round buffer, which are already
// padded and shuffled, so this can be done non-obliviously.  tot_msgs
// is the total number of messages in the input buffer, which may
// include padding messages added by the shuffle.  Those messages are
// not sent anywhere.  There are num_msgs_per_stg messages for each
// storage node labelled for that node.
static void send_round2_msgs(uint32_t tot_msgs, uint32_t num_msgs_per_stg, MsgBuffer &prevround)
{
    uint16_t msg_size = g_teems_config.msg_size;
    const uint8_t* buf = prevround.buf;
    nodenum_t num_storage_nodes = g_teems_config.num_storage_nodes;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    uint8_t *myself_buf = NULL;

    for (nodenum_t i=0; i<num_storage_nodes; ++i) {
        nodenum_t node = g_teems_config.storage_nodes[i];
        if (node != my_node_num) {
            g_commstates[node].message_start(msg_size * num_msgs_per_stg);
        } else {
            MsgBuffer &round2 = route_state.round2;
            pthread_mutex_lock(&round2.mutex);
            uint32_t start = round2.reserved;
            if (start + num_msgs_per_stg > round2.bufsize) {
                pthread_mutex_unlock(&round2.mutex);
                printf("Max %u messages exceeded\n", round2.bufsize);
                return;
            }
            round2.reserved += num_msgs_per_stg;
            pthread_mutex_unlock(&round2.mutex);
            myself_buf = round2.buf + start * msg_size;
        }
    }

    while (tot_msgs) {
        nodenum_t storage_node_id =
            nodenum_t((*(const uint32_t *)buf)>>DEST_UID_BITS);
        if (storage_node_id < num_storage_nodes) {
            nodenum_t node = g_teems_config.storage_map[storage_node_id];
            if (node == my_node_num) {
                memmove(myself_buf, buf, msg_size);
                myself_buf += msg_size;
            } else {
                g_commstates[node].message_data(buf, msg_size);
            }
        }

        buf += msg_size;
        --tot_msgs;
    }

    if (myself_buf) {
        MsgBuffer &round2 = route_state.round2;
        pthread_mutex_lock(&round2.mutex);
        round2.inserted += num_msgs_per_stg;
        round2.nodes_received += 1;
        pthread_mutex_unlock(&round2.mutex);
    }
}


static void round1a_processing(void *cbpointer) {
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];
    MsgBuffer &round1 = route_state.round1;

    if (my_roles & ROLE_ROUTING) {
        route_state.cbpointer = cbpointer;
        pthread_mutex_lock(&round1.mutex);
        // Ensure there are no pending messages currently being inserted
        // into the buffer
        while (round1.reserved != round1.inserted) {
            pthread_mutex_unlock(&round1.mutex);
            pthread_mutex_lock(&round1.mutex);
        }
#ifdef TRACE_ROUTING
        show_messages("Start of round 1a", round1.buf, round1.inserted);
#endif
#ifdef PROFILE_ROUTING
        uint32_t inserted = round1.inserted;
        unsigned long start_round1a =
            printf_with_rtclock("begin round1a processing (%u)\n", inserted);
        // Sort the messages we've received
        unsigned long start_sort =
            printf_with_rtclock("begin oblivious sort (%u,%u)\n", inserted,
                route_state.max_round1_msgs);
#endif
        // Sort received messages by increasing user ID and
        // priority. Larger priority number indicates higher priority.
        sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, round1.buf,
            g_teems_config.msg_size, round1.inserted, route_state.max_round1_msgs,
            send_round1a_msgs);

#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_sort, "end oblivious sort (%u,%u)\n",
            inserted, route_state.max_round1_msgs);
        printf_with_rtclock_diff(start_round1a, "end round1a processing (%u)\n",
            inserted);
#endif
        round1.reset();
        pthread_mutex_unlock(&round1.mutex);

        MsgBuffer &round1a = route_state.round1a;
        pthread_mutex_lock(&round1a.mutex);
        round1a.completed_prev_round = true;
        nodenum_t nodes_received = round1a.nodes_received;
        pthread_mutex_unlock(&round1a.mutex);
        if (nodes_received == g_teems_config.num_routing_nodes) {
            route_state.step = ROUTE_ROUND_1A;
            route_state.cbpointer = NULL;
            ocall_routing_round_complete(cbpointer, ROUND_1A);
        }
    } else {
        route_state.step = ROUTE_ROUND_1A;
        route_state.round1a.completed_prev_round = true;
        ocall_routing_round_complete(cbpointer, ROUND_1A);
    }
}

static void round1b_processing(void *cbpointer) {
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];
    nodenum_t my_node_num = g_teems_config.my_node_num;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    MsgBuffer &round1a = route_state.round1a;
    MsgBuffer &round1a_sorted = route_state.round1a_sorted;

    if (my_roles & ROLE_ROUTING) {
        route_state.cbpointer = cbpointer;
        pthread_mutex_lock(&round1a.mutex);
        // Ensure there are no pending messages currently being inserted
        // into the buffer
        while (round1a.reserved != round1a.inserted) {
            pthread_mutex_unlock(&round1a.mutex);
            pthread_mutex_lock(&round1a.mutex);
        }
        pthread_mutex_lock(&round1a_sorted.mutex);

#ifdef TRACE_ROUTING
        show_messages("Start of round 1b", round1a.buf, round1a.inserted);
#endif
#ifdef PROFILE_ROUTING
        uint32_t inserted = round1a.inserted;
        unsigned long start_round1b =
            printf_with_rtclock("begin round1b processing (%u)\n", inserted);
        // Sort the messages we've received
#endif
        // Sort received messages by increasing user ID and
        // priority. Larger priority number indicates higher priority.
        if (inserted > 0) {
            // copy items in sorted order into round1a_sorted
#ifdef PROFILE_ROUTING
        unsigned long start_sort =
            printf_with_rtclock("begin oblivious sort (%u,%u)\n", inserted,
                route_state.max_round1a_msgs);
#endif
            sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, round1a.buf,
                g_teems_config.msg_size, round1a.inserted,
                route_state.max_round1a_msgs,
                route_state.round1a_sorted.buf);
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_sort, "end oblivious sort (%u,%u)\n",
            inserted, route_state.max_round1a_msgs);
#endif
            send_round1b_msgs(round1a_sorted.buf, round1a.inserted);
        } else {
            send_round1b_msgs(NULL, 0);
        }

#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_round1b, "end round1b processing (%u)\n",
            inserted);
#endif

        pthread_mutex_unlock(&round1a_sorted.mutex);
        pthread_mutex_unlock(&round1a.mutex);
        MsgBuffer &round1b_prev = route_state.round1b_prev;
        pthread_mutex_lock(&round1b_prev.mutex);
        round1b_prev.completed_prev_round = true;
        nodenum_t nodes_received = round1b_prev.nodes_received;
        pthread_mutex_unlock(&round1b_prev.mutex);
        MsgBuffer &round1b_next = route_state.round1b_next;
        pthread_mutex_lock(&round1b_next.mutex);
        round1b_next.completed_prev_round = true;
        nodes_received += round1b_next.nodes_received;
        pthread_mutex_unlock(&round1b_next.mutex);
        nodenum_t adjacent_nodes;
        if (num_routing_nodes == 1) {
            adjacent_nodes = 0;
        } else if ((prev_nodes == 0) || (prev_nodes == num_routing_nodes-1)) {
            adjacent_nodes = 1;
        } else {
            adjacent_nodes = 2;
        }
        if (nodes_received == adjacent_nodes) {
            route_state.step = ROUTE_ROUND_1B;
            route_state.cbpointer = NULL;
            ocall_routing_round_complete(cbpointer, ROUND_1B);
        }
    } else {
        route_state.step = ROUTE_ROUND_1B;
        route_state.round1b_prev.completed_prev_round = true;
        route_state.round1b_next.completed_prev_round = true;
        ocall_routing_round_complete(cbpointer, ROUND_1B);
    }
}

static void copy_msgs(uint8_t *dst, uint32_t start_msg, uint32_t num_copy, const uint8_t *src,
    const UidPriorityKey *indices)
{
    uint16_t msg_size = g_teems_config.msg_size;
    const UidPriorityKey *idxp = indices + start_msg;
    uint8_t *buf = dst;
    for (uint32_t i=0; i<num_copy; i++) {
        memmove(buf, src + idxp[i].index()*msg_size, msg_size);
        buf += msg_size;
    }
}

static void round1c_processing(void *cbpointer) {
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];
    nodenum_t my_node_num = g_teems_config.my_node_num;
    nodenum_t num_routing_nodes = g_teems_config.num_routing_nodes;
    uint16_t prev_nodes = g_teems_config.weights[my_node_num].startweight;
    uint16_t msg_size = g_teems_config.msg_size;
    uint32_t max_round1b_msgs_to_adj_rtr = route_state.max_round1b_msgs_to_adj_rtr;
    uint32_t max_round1a_msgs = route_state.max_round1a_msgs;
    MsgBuffer &round1a = route_state.round1a;
    MsgBuffer &round1a_sorted = route_state.round1a_sorted;
    MsgBuffer &round1b_prev = route_state.round1b_prev;
    MsgBuffer &round1b_next = route_state.round1b_next;

    if (my_roles & ROLE_ROUTING) {
        route_state.cbpointer = cbpointer;
        pthread_mutex_lock(&round1b_prev.mutex);
        pthread_mutex_lock(&round1b_next.mutex);
        // Ensure there are no pending messages currently being inserted
        // into the round 1b buffers
        while (round1b_prev.reserved != round1b_prev.inserted) {
            pthread_mutex_unlock(&round1b_prev.mutex);
            pthread_mutex_lock(&round1b_prev.mutex);
        }
        while (round1b_next.reserved != round1b_next.inserted) {
            pthread_mutex_unlock(&round1b_next.mutex);
            pthread_mutex_lock(&round1b_next.mutex);
        }
        pthread_mutex_lock(&round1a.mutex);
        pthread_mutex_lock(&round1a_sorted.mutex);

#ifdef TRACE_ROUTING
        show_messages("Start of round 1c", round1a.buf, round1a.inserted);
#endif
#ifdef PROFILE_ROUTING
        unsigned long start_round1c = printf_with_rtclock("begin round1c processing (%u)\n", round1a.inserted);
#endif

        // sort round1b_prev msgs with initial msgs in round1a_sorted
        if (prev_nodes > 0) {
            // Copy initial msgs in round1a_sorted to round1b_prev buffer for sorting
            // Note that all inserted values and buffer sizes are non-secret
            uint32_t num_init_round1a = min(round1a.inserted,
                max_round1b_msgs_to_adj_rtr);
            uint32_t num_round1b_prev = round1b_prev.inserted;
            if (num_round1b_prev + num_init_round1a  <= max_round1b_msgs_to_adj_rtr) {
                // all our round 1a messages "belong" to previous router and can be removed here
                round1a.inserted = 0;
            } else {
                // copy initial round1a msgs after round1b_prev msgs
                memmove(round1b_prev.buf+num_round1b_prev*msg_size, round1a_sorted.buf,
                    num_init_round1a*msg_size);
                // sort and take final msgs as initial round1a msgs
#ifdef PROFILE_ROUTING
                unsigned long start_sort = printf_with_rtclock("begin round1b_prev oblivious sort (%u,%u)\n", num_round1b_prev + num_init_round1a, 2*max_round1b_msgs_to_adj_rtr);
#endif
                uint32_t num_copy = num_round1b_prev+num_init_round1a-max_round1b_msgs_to_adj_rtr;
                sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, round1b_prev.buf,
                    msg_size, num_round1b_prev + num_init_round1a,
                    2*max_round1b_msgs_to_adj_rtr,
                    [&](const uint8_t *src, const UidPriorityKey *indices, uint32_t Nr) {
                        return copy_msgs(round1a_sorted.buf, max_round1b_msgs_to_adj_rtr,
                            num_copy, src, indices);
                    }
                );
                round1a.inserted -= (max_round1b_msgs_to_adj_rtr-num_round1b_prev);
#ifdef PROFILE_ROUTING
                printf_with_rtclock_diff(start_sort, "end round1b_prev oblivious sort (%u,%u)\n", num_round1b_prev + num_init_round1a, 2*max_round1b_msgs_to_adj_rtr);
#endif
            }
        }
        // sort round1b_next msgs with final msgs in round1a_sorted
        if ((prev_nodes < num_routing_nodes-1) && (round1b_next.inserted > 0)) {
            // Copy final msgs in round1a_sorted to round1b_next buffer for sorting
            // Note that all inserted values and buffer sizes are non-secret
            // round1b_next.inserted>0, so round1a >= max_round1a_msgs-max_round1b_msgs_to_adj_rtr
            uint32_t round1a_msg_start = max_round1a_msgs-max_round1b_msgs_to_adj_rtr;
            uint32_t num_final_round1a = round1a.inserted - round1a_msg_start;
            uint32_t num_round1b_next = round1b_next.inserted;
            memmove(round1b_next.buf+num_round1b_next*msg_size,
                round1a_sorted.buf + round1a_msg_start*msg_size,
                num_final_round1a*msg_size);
            // sort and take initial msgs as final round1a msgs
#ifdef PROFILE_ROUTING
            unsigned long start_sort = printf_with_rtclock("begin round1b_next oblivious sort (%u,%u)\n", num_round1b_next + num_final_round1a, 2*max_round1b_msgs_to_adj_rtr);
#endif
            uint32_t num_copy = min(num_final_round1a+num_round1b_next,
                max_round1b_msgs_to_adj_rtr);
            sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, round1b_next.buf,
                msg_size, num_round1b_next + num_final_round1a, 2*max_round1b_msgs_to_adj_rtr,
                [&](const uint8_t *src, const UidPriorityKey *indices, uint32_t Nr) {
                    return copy_msgs(round1a_sorted.buf + round1a_msg_start*msg_size, 0,
                        num_copy, src, indices);
                }
            );
            round1a.inserted += (num_copy - num_final_round1a);
#ifdef PROFILE_ROUTING
            printf_with_rtclock_diff(start_sort, "end round1b_next oblivious sort (%u,%u)\n", num_round1b_next + num_final_round1a, 2*max_round1b_msgs_to_adj_rtr);
#endif
        }
#ifdef TRACE_ROUTING
        printf("round1a_sorted.inserted = %lu. round1a.inserted = %lu\n",
            round1a_sorted.inserted, round1a.inserted);
        show_messages("In round 1c", round1a_sorted.buf,
            round1a.inserted);
#endif
#ifdef PROFILE_ROUTING
        unsigned long start_sort = printf_with_rtclock("begin full oblivious sort (%u,%u)\n", round1a.inserted, route_state.max_round1a_msgs);
#endif
        // Sort received messages by increasing user ID and
        // priority. Larger priority number indicates higher priority.
        if (round1a.inserted > 0) {
            sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, round1a_sorted.buf,
                msg_size, round1a.inserted, route_state.max_round1a_msgs,
                [&](const uint8_t *msgs, const UidPriorityKey *indices, uint32_t N) {
                    send_round_robin_msgs<UidPriorityKey>(route_state.round1c, msgs, indices, N);
                });
        } else {
            send_round_robin_msgs<UidPriorityKey>(route_state.round1c, NULL, NULL, 0);
        }

#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_sort, "end full oblivious sort (%u,%u)\n", round1a.inserted, route_state.max_round1a_msgs);
        printf_with_rtclock_diff(start_round1c, "end round1c processing (%u)\n", round1a.inserted);
#endif

        round1a.reset();
        round1a_sorted.reset();
        round1b_prev.reset();
        round1b_next.reset();
        pthread_mutex_unlock(&round1a_sorted.mutex);
        pthread_mutex_unlock(&round1a.mutex);
        pthread_mutex_unlock(&round1b_next.mutex);
        pthread_mutex_unlock(&round1b_prev.mutex);

        MsgBuffer &round1c = route_state.round1c;
        pthread_mutex_lock(&round1c.mutex);
        round1c.completed_prev_round = true;
        nodenum_t nodes_received = round1c.nodes_received;
        pthread_mutex_unlock(&round1c.mutex);
        if (nodes_received == num_routing_nodes) {
            route_state.step = ROUTE_ROUND_1C;
            route_state.cbpointer = NULL;
            ocall_routing_round_complete(cbpointer, ROUND_1C);
        }
    } else {
        route_state.step = ROUTE_ROUND_1C;
        route_state.round1c.completed_prev_round = true;
        ocall_routing_round_complete(cbpointer, ROUND_1C);
    }
}

// Process messages in round 2
static void round2_processing(uint8_t my_roles, void *cbpointer, MsgBuffer &prevround) {
    if (my_roles & ROLE_ROUTING) {
        route_state.cbpointer = cbpointer;

        pthread_mutex_lock(&prevround.mutex);
        // Ensure there are no pending messages currently being inserted
        // into the buffer
        while (prevround.reserved != prevround.inserted) {
            pthread_mutex_unlock(&prevround.mutex);
            pthread_mutex_lock(&prevround.mutex);
        }

        // If the _total_ number of messages we received in round 1
        // is less than the max number of messages we could send to
        // _each_ storage node, then cap the number of messages we
        // will send to each storage node to that number.
        uint32_t msgs_per_stg = route_state.max_msg_to_each_stg;
        if (prevround.inserted < msgs_per_stg) {
            msgs_per_stg = prevround.inserted;
        }

        // Note: at this point, it is required that each message in
        // the prevround buffer have a _valid_ storage node id field.

        // Obliviously tally the number of messages we received in
        // the previous round destined for each storage node
#ifdef TRACE_ROUTING
        show_messages("Start of round 2", prevround.buf, prevround.inserted);
#endif
#ifdef PROFILE_ROUTING
        unsigned long start_round2 = printf_with_rtclock("begin round2 processing (%u,%u)\n", prevround.inserted, prevround.bufsize);
        unsigned long start_tally = printf_with_rtclock("begin tally (%u)\n", prevround.inserted);
#endif
        uint16_t msg_size = g_teems_config.msg_size;
        nodenum_t num_storage_nodes = g_teems_config.num_storage_nodes;
        std::vector<uint32_t> tally = obliv_tally_stg(
            prevround.buf, msg_size, prevround.inserted, num_storage_nodes);
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_tally, "end tally (%u)\n", prevround.inserted);
#endif
        // Note: tally contains private values!  It's OK to
        // non-obliviously check for an error condition, though.
        // While we're at it, obliviously change the tally of
        // messages received to a tally of padding messages
        // required.
        uint32_t tot_padding = 0;
        for (nodenum_t i=0; i<num_storage_nodes; ++i) {
            if (tally[i] > msgs_per_stg) {
                printf("Received too many messages for storage node %u\n", i);
                assert(tally[i] <= msgs_per_stg);
            }
            tally[i] = msgs_per_stg - tally[i];
            tot_padding += tally[i];
        }

        prevround.reserved += tot_padding;
        assert(prevround.reserved <= prevround.bufsize);

        // Obliviously add padding for each storage node according
        // to the (private) padding tally.
#ifdef PROFILE_ROUTING
        unsigned long start_pad = printf_with_rtclock("begin pad (%u)\n", tot_padding);
#endif
        obliv_pad_stg(prevround.buf + prevround.inserted * msg_size,
            msg_size, tally, tot_padding);
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_pad, "end pad (%u)\n", tot_padding);
#endif

        prevround.inserted += tot_padding;

        // Obliviously shuffle the messages
#ifdef PROFILE_ROUTING
        unsigned long start_shuffle = printf_with_rtclock("begin shuffle (%u,%u)\n", prevround.inserted, prevround.bufsize);
#endif
        uint32_t num_shuffled = shuffle_mtobliv(g_teems_config.nthreads,
            prevround.buf, msg_size, prevround.inserted, prevround.bufsize);
#ifdef PROFILE_ROUTING
        printf_with_rtclock_diff(start_shuffle, "end shuffle (%u,%u)\n", prevround.inserted, prevround.bufsize);
        printf_with_rtclock_diff(start_round2, "end round2 processing (%u,%u)\n", prevround.inserted, prevround.bufsize);
#endif

        // Now we can handle the messages non-obliviously, since we
        // know there will be exactly msgs_per_stg messages to each
        // storage node, and the oblivious shuffle broke the
        // connection between where each message came from and where
        // it's going.
        send_round2_msgs(num_shuffled, msgs_per_stg, prevround);

        prevround.reset();
        pthread_mutex_unlock(&prevround.mutex);
    }

    if (my_roles & ROLE_STORAGE) {
        route_state.cbpointer = cbpointer;
        MsgBuffer &round2 = route_state.round2;

        pthread_mutex_lock(&round2.mutex);
        round2.completed_prev_round = true;
        nodenum_t nodes_received = round2.nodes_received;
        pthread_mutex_unlock(&round2.mutex);

        if (nodes_received == g_teems_config.num_routing_nodes) {
            route_state.step = ROUTE_ROUND_2;
            route_state.cbpointer = NULL;
            ocall_routing_round_complete(cbpointer, 2);
        }
    } else {
        route_state.step = ROUTE_ROUND_2;
        route_state.round2.completed_prev_round = true;
        ocall_routing_round_complete(cbpointer, 2);
    }
}

// Perform the next round of routing.  The callback pointer will be
// passed to ocall_routing_round_complete when the round is complete.
void ecall_routing_proceed(void *cbpointer)
{
    uint8_t my_roles = g_teems_config.roles[g_teems_config.my_node_num];

    if (route_state.step == ROUTE_NOT_STARTED) {
        if (my_roles & ROLE_INGESTION) {
            ingestion_epoch++;
            route_state.cbpointer = cbpointer;
            MsgBuffer &ingbuf = route_state.ingbuf;

            pthread_mutex_lock(&ingbuf.mutex);
            // Ensure there are no pending messages currently being inserted
            // into the buffer
            while (ingbuf.reserved != ingbuf.inserted) {
                pthread_mutex_unlock(&ingbuf.mutex);
                pthread_mutex_lock(&ingbuf.mutex);
            }
            // Sort the messages we've received
#ifdef TRACE_ROUTING
            show_messages("Start of round 1", ingbuf.buf, ingbuf.inserted);
#endif
#ifdef PROFILE_ROUTING
            uint32_t inserted = ingbuf.inserted;
            unsigned long start_round1 =
                printf_with_rtclock("begin round1 processing (%u)\n", inserted);
            unsigned long start_sort =
                printf_with_rtclock("begin oblivious sort (%u,%u)\n", inserted,
                    route_state.tot_msg_per_ing);
#endif
            if (g_teems_config.private_routing) {
                sort_mtobliv<UidKey>(g_teems_config.nthreads, ingbuf.buf,
                    g_teems_config.msg_size, ingbuf.inserted,
                    route_state.tot_msg_per_ing,
                    [&](const uint8_t *msgs, const UidKey *indices, uint32_t N) {
                        send_round_robin_msgs<UidKey>(route_state.round1,
                            msgs, indices, N);
                    });
            } else {
                // Sort received messages by increasing user ID and
                // priority. Larger priority number indicates higher priority.
                sort_mtobliv<UidPriorityKey>(g_teems_config.nthreads, ingbuf.buf,
                    g_teems_config.msg_size, ingbuf.inserted,
                    route_state.tot_msg_per_ing,
                    [&](const uint8_t *msgs,
                            const UidPriorityKey *indices, uint32_t N) {
                        send_round_robin_msgs<UidPriorityKey>(route_state.round1,
                            msgs, indices, N);
                    });
            }
#ifdef PROFILE_ROUTING
            printf_with_rtclock_diff(start_sort, "end oblivious sort (%u,%u)\n",
                inserted, route_state.tot_msg_per_ing);
            printf_with_rtclock_diff(start_round1, "end round1 processing (%u)\n",
                inserted);
#endif
            ingbuf.reset();
            pthread_mutex_unlock(&ingbuf.mutex);
        }
        if (my_roles & ROLE_ROUTING) {
            MsgBuffer &round1 = route_state.round1;

            pthread_mutex_lock(&round1.mutex);
            round1.completed_prev_round = true;
            nodenum_t nodes_received = round1.nodes_received;
            pthread_mutex_unlock(&round1.mutex);

            if (nodes_received == g_teems_config.num_ingestion_nodes) {
                route_state.step = ROUTE_ROUND_1;
                route_state.cbpointer = NULL;
                ocall_routing_round_complete(cbpointer, 1);
            }
        } else {
            route_state.step = ROUTE_ROUND_1;
            route_state.round1.completed_prev_round = true;
            ocall_routing_round_complete(cbpointer, 1);
        }
    } else if (route_state.step == ROUTE_ROUND_1) {
        if (g_teems_config.private_routing) { // private routing next round
            round2_processing(my_roles, cbpointer, route_state.round1);
        } else { // public routing next round
            round1a_processing(cbpointer);
        }
    } else if (route_state.step == ROUTE_ROUND_1A) {
        round1b_processing(cbpointer);
    } else if (route_state.step == ROUTE_ROUND_1B) {
        round1c_processing(cbpointer);
    } else if (route_state.step == ROUTE_ROUND_1C) {
        round2_processing(my_roles, cbpointer, route_state.round1c);
    } else if (route_state.step == ROUTE_ROUND_2) {
        if (my_roles & ROLE_STORAGE) {
            MsgBuffer &round2 = route_state.round2;

            pthread_mutex_lock(&round2.mutex);
            // Ensure there are no pending messages currently being inserted
            // into the buffer
            while (round2.reserved != round2.inserted) {
                pthread_mutex_unlock(&round2.mutex);
                pthread_mutex_lock(&round2.mutex);
            }

            unsigned long start = printf_with_rtclock("begin storage processing (%u)\n", round2.inserted);
            storage_received(round2);
            printf_with_rtclock_diff(start, "end storage processing (%u)\n", round2.inserted);

            // We're done
            route_state.step = ROUTE_NOT_STARTED;
            ocall_routing_round_complete(cbpointer, 0);
        } else {
            // We're done
            route_state.step = ROUTE_NOT_STARTED;
            ocall_routing_round_complete(cbpointer, 0);
        }
    }
}