Integrate credis with Ray & route task table entries into credis. (#1841)

This commit is contained in:
Zongheng Yang
2018-05-24 23:35:25 -07:00
committed by Robert Nishihara
parent 3509a33cf3
commit fa97acbc89
25 changed files with 698 additions and 352 deletions
+61
View File
@@ -0,0 +1,61 @@
#ifndef RAY_CHAIN_MODULE_H_
#define RAY_CHAIN_MODULE_H_
#include <functional>
#include "redismodule.h"
// NOTE(zongheng): this duplicated declaration serves as forward-declaration
// only. The implementation is supposed to be linked in from credis. In
// principle, we can expose a header from credis and simple include that header.
// This is left as future work.
//
// Concrete definitions from credis (from an example commit):
// https://github.com/ray-project/credis/blob/7eae7f2e58d16dfa1a95b5dfab02549f54b94e5d/src/member.cc#L41
// https://github.com/ray-project/credis/blob/7eae7f2e58d16dfa1a95b5dfab02549f54b94e5d/src/master.cc#L36
// Typical usage to make an existing redismodule command chain-compatible:
//
// extern RedisChainModule module;
// int MyCmd_RedisModuleCmd(...) {
// return module.Mutate(..., NodeFunc, TailFunc);
// }
//
// See, for instance, ChainTableAdd_RedisCommand in ray_redis_module.cc.
class RedisChainModule {
public:
// A function that runs on every node in the chain. Type:
// (context, argv, argc, (can be nullptr) mutated_key_str) -> int
//
// (Advanced) The optional fourth arg can be used in the following way:
//
// RedisModuleString* redis_key_str = nullptr;
// node_func(ctx, argv, argc, &redis_key_str);
// // "redis_key_str" now points to the RedisModuleString whose contents
// // is mutated by "node_func".
//
// If the fourth arg is passed, NodeFunc *must* fill in the key being mutated.
// It is okay for this NodeFunc to call "RM_FreeString(mutated_key_str)" after
// assigning the fourth arg, since that call presumably only decrements a ref
// count.
using NodeFunc = std::function<
int(RedisModuleCtx *, RedisModuleString **, int, RedisModuleString **)>;
// A function that (1) runs only after all NodeFunc's have run, and (2) runs
// once on the tail. A typical usage is to publish a write.
using TailFunc =
std::function<int(RedisModuleCtx *, RedisModuleString **, int)>;
// TODO(zongheng): document the RM_Reply semantics.
// Runs "node_func" on every node in the chain; after the tail node has run it
// too, finalizes the mutation by running "tail_func".
// TODO(zongheng): currently only supports 1-node chain.
int ChainReplicate(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc,
NodeFunc node_func,
TailFunc tail_func);
};
#endif // RAY_CHAIN_MODULE_H_
+114 -39
View File
@@ -1,14 +1,26 @@
#include "redismodule.h"
#include <string.h>
#include "redis_string.h"
#include "common_protocol.h"
#include "format/common_generated.h"
#include "ray/gcs/format/gcs_generated.h"
#include "ray/id.h"
#include "redis_string.h"
#include "redismodule.h"
#include "task.h"
#if RAY_USE_NEW_GCS
// Under this flag, ray-project/credis will be loaded. Specifically, via
// "path/redis-server --loadmodule <credis module> --loadmodule <current
// libray_redis_module>" (dlopen() under the hood) will a definition of "module"
// be supplied.
//
// All commands in this file that depend on "module" must be wrapped by "#if
// RAY_USE_NEW_GCS", until we switch to this launch configuration as the
// default.
#include "chain_module.h"
extern RedisChainModule module;
#endif
// Various tables are maintained in redis:
//
// == OBJECT TABLE ==
@@ -88,9 +100,14 @@ RedisModuleString *FormatPubsubChannel(
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
const char *prefix,
RedisModuleString *keyname,
int mode) {
int mode,
RedisModuleString **mutated_key_str) {
RedisModuleString *prefixed_keyname =
RedisString_Format(ctx, "%s%S", prefix, keyname);
// Pass out the key being mutated, should the caller request so.
if (mutated_key_str != nullptr) {
*mutated_key_str = prefixed_keyname;
}
RedisModuleKey *key =
(RedisModuleKey *) RedisModule_OpenKey(ctx, prefixed_keyname, mode);
return key;
@@ -99,7 +116,8 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
RedisModuleString *prefix_enum,
RedisModuleString *keyname,
int mode) {
int mode,
RedisModuleString **mutated_key_str) {
long long prefix_long;
RAY_CHECK(RedisModule_StringToLongLong(prefix_enum, &prefix_long) ==
REDISMODULE_OK)
@@ -109,7 +127,24 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
<< "This table has no prefix registered";
RAY_CHECK(prefix >= TablePrefix_MIN && prefix <= TablePrefix_MAX)
<< "Prefix must be a valid TablePrefix";
return OpenPrefixedKey(ctx, table_prefixes[prefix], keyname, mode);
return OpenPrefixedKey(ctx, table_prefixes[prefix], keyname, mode,
mutated_key_str);
}
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
const char *prefix,
RedisModuleString *keyname,
int mode) {
return OpenPrefixedKey(ctx, prefix, keyname, mode,
/*mutated_key_str=*/nullptr);
}
RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
RedisModuleString *prefix_enum,
RedisModuleString *keyname,
int mode) {
return OpenPrefixedKey(ctx, prefix_enum, keyname, mode,
/*mutated_key_str=*/nullptr);
}
/// Open the key used to store the channels that should be published to when an
@@ -444,11 +479,12 @@ bool PublishObjectNotification(RedisModuleCtx *ctx,
// NOTE(pcmoritz): This is a temporary redis command that will be removed once
// the GCS uses https://github.com/pcmoritz/credis.
int TaskTableAdd(RedisModuleCtx *ctx,
RedisModuleString *id,
RedisModuleString *data) {
int PublishTaskTableAdd(RedisModuleCtx *ctx,
RedisModuleString *id,
RedisModuleString *data) {
const char *buf = RedisModule_StringPtrLen(data, NULL);
auto message = flatbuffers::GetRoot<TaskTableData>(buf);
RAY_CHECK(message != nullptr);
if (message->scheduling_state() == SchedulingState_WAITING ||
message->scheduling_state() == SchedulingState_SCHEDULED) {
@@ -483,7 +519,6 @@ int TaskTableAdd(RedisModuleCtx *ctx,
long long num_clients = RedisModule_CallReplyInteger(reply);
RAY_CHECK(num_clients <= 1) << "Published to " << num_clients
<< " clients.";
}
return RedisModule_ReplyWithSimpleString(ctx, "OK");
}
@@ -542,47 +577,47 @@ int PublishTableAdd(RedisModuleCtx *ctx,
return RedisModule_ReplyWithSimpleString(ctx, "OK");
}
/// Add an entry at a key. This overwrites any existing data at the key.
/// Publishes a notification about the update to all subscribers, if a pubsub
/// channel is provided.
///
/// This is called from a client with the command:
//
/// RAY.TABLE_ADD <table_prefix> <pubsub_channel> <id> <data>
///
/// \param table_prefix The prefix string for keys in this table.
/// \param pubsub_channel The pubsub channel name that notifications for
/// this key should be published to. When publishing to a specific
/// client, the channel name should be <pubsub_channel>:<client_id>.
/// \param id The ID of the key to set.
/// \param data The data to insert at the key.
/// \return The current value at the key, or OK if there is no value.
int TableAdd_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
RedisModule_AutoMemory(ctx);
// RAY.TABLE_ADD:
// TableAdd_RedisCommand: the actual command handler.
// (helper) TableAdd_DoWrite: performs the write to redis state.
// (helper) TableAdd_DoPublish: performs a publish after the write.
// ChainTableAdd_RedisCommand: the same command, chain-enabled.
int TableAdd_DoWrite(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc,
RedisModuleString **mutated_key_str) {
if (argc != 5) {
return RedisModule_WrongArity(ctx);
}
RedisModuleString *prefix_str = argv[1];
RedisModuleString *id = argv[3];
RedisModuleString *data = argv[4];
RedisModuleKey *key =
OpenPrefixedKey(ctx, prefix_str, id, REDISMODULE_READ | REDISMODULE_WRITE,
mutated_key_str);
RedisModule_StringSet(key, data);
return REDISMODULE_OK;
}
int TableAdd_DoPublish(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
if (argc != 5) {
return RedisModule_WrongArity(ctx);
}
RedisModuleString *pubsub_channel_str = argv[2];
RedisModuleString *id = argv[3];
RedisModuleString *data = argv[4];
// Set the keys in the table.
RedisModuleKey *key = OpenPrefixedKey(ctx, prefix_str, id,
REDISMODULE_READ | REDISMODULE_WRITE);
RedisModule_StringSet(key, data);
// Publish a message on the requested pubsub channel if necessary.
TablePubsub pubsub_channel = ParseTablePubsub(pubsub_channel_str);
if (pubsub_channel == TablePubsub_TASK) {
// Publish the task to its subscribers.
// TODO(swang): This is only necessary for legacy Ray and should be removed
// once we switch to using the new GCS API for the task table.
return TaskTableAdd(ctx, id, data);
return PublishTaskTableAdd(ctx, id, data);
} else if (pubsub_channel != TablePubsub_NO_PUBLISH) {
// All other pubsub channels write the data back directly onto the channel.
return PublishTableAdd(ctx, pubsub_channel_str, id, data);
@@ -591,6 +626,38 @@ int TableAdd_RedisCommand(RedisModuleCtx *ctx,
}
}
/// Add an entry at a key. This overwrites any existing data at the key.
/// Publishes a notification about the update to all subscribers, if a pubsub
/// channel is provided.
///
/// This is called from a client with the command:
///
/// RAY.TABLE_ADD <table_prefix> <pubsub_channel> <id> <data>
///
/// \param table_prefix The prefix string for keys in this table.
/// \param pubsub_channel The pubsub channel name that notifications for
/// this key should be published to. When publishing to a specific
/// client, the channel name should be <pubsub_channel>:<client_id>.
/// \param id The ID of the key to set.
/// \param data The data to insert at the key.
/// \return The current value at the key, or OK if there is no value.
int TableAdd_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
RedisModule_AutoMemory(ctx);
TableAdd_DoWrite(ctx, argv, argc, /*mutated_key_str=*/nullptr);
return TableAdd_DoPublish(ctx, argv, argc);
}
#if RAY_USE_NEW_GCS
int ChainTableAdd_RedisCommand(RedisModuleCtx *ctx,
RedisModuleString **argv,
int argc) {
return module.ChainReplicate(ctx, argv, argc, /*node_func=*/TableAdd_DoWrite,
/*tail_func=*/TableAdd_DoPublish);
}
#endif
/// Append an entry to the log stored at a key. Publishes a notification about
/// the update to all subscribers, if a pubsub channel is provided.
///
@@ -746,7 +813,6 @@ int TableLookup_RedisCommand(RedisModuleCtx *ctx,
ctx, reinterpret_cast<const char *>(fbb.GetBufferPointer()),
fbb.GetSize());
}
return REDISMODULE_OK;
}
@@ -1663,7 +1729,7 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
}
if (RedisModule_CreateCommand(ctx, "ray.table_add", TableAdd_RedisCommand,
"write", 0, 0, 0) == REDISMODULE_ERR) {
"write pubsub", 0, 0, 0) == REDISMODULE_ERR) {
return REDISMODULE_ERR;
}
@@ -1763,6 +1829,15 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
return REDISMODULE_ERR;
}
#if RAY_USE_NEW_GCS
// Chain-enabled commands that depend on ray-project/credis.
if (RedisModule_CreateCommand(ctx, "ray.chain.table_add",
ChainTableAdd_RedisCommand, "write pubsub", 0,
0, 0) == REDISMODULE_ERR) {
return REDISMODULE_ERR;
}
#endif
return REDISMODULE_OK;
}
+8 -1
View File
@@ -1,6 +1,11 @@
#ifndef RAY_REDIS_STRING_H_
#define RAY_REDIS_STRING_H_
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include "redismodule.h"
/* Format a RedisModuleString.
*
@@ -65,3 +70,5 @@ RedisModuleString *RedisString_Format(RedisModuleCtx *ctx,
va_end(ap);
return result;
}
#endif // RAY_REDIS_STRING_H_
+21 -4
View File
@@ -3,12 +3,29 @@
# This needs to be run in the build tree, which is normally ray/build
# Cause the script to exit if a single command fails.
set -e
set -ex
LaunchRedis() {
port=$1
if [[ "${RAY_USE_NEW_GCS}" = "on" ]]; then
./src/credis/redis/src/redis-server \
--loglevel warning \
--loadmodule ./src/credis/build/src/libmember.so \
--loadmodule ./src/common/redis_module/libray_redis_module.so \
--port $port &
else
./src/common/thirdparty/redis/src/redis-server \
--loglevel warning \
--loadmodule ./src/common/redis_module/libray_redis_module.so \
--port $port &
fi
sleep 1s
}
# Start the Redis shards.
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
sleep 1s
LaunchRedis 6379
LaunchRedis 6380
# Register the shard location with the primary shard.
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
+10 -11
View File
@@ -7,22 +7,21 @@ set -x
# Cause the script to exit if a single command fails.
set -e
# Start the Redis shards.
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
sleep 1s
# Register the shard location with the primary shard.
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
if [ -z "$RAY_USE_NEW_GCS" ]; then
# Start the Redis shards.
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6379 &
./src/common/thirdparty/redis/src/redis-server --loglevel warning --loadmodule ./src/common/redis_module/libray_redis_module.so --port 6380 &
sleep 1s
# Register the shard location with the primary shard.
./src/common/thirdparty/redis/src/redis-cli set NumRedisShards 1
./src/common/thirdparty/redis/src/redis-cli rpush RedisShards 127.0.0.1:6380
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/db_tests
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/io_tests
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_tests
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/redis_tests
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/task_table_tests
valgrind --track-origins=yes --leak-check=full --show-leak-kinds=all --leak-check-heuristics=stdstring --error-exitcode=1 ./src/common/object_table_tests
./src/common/thirdparty/redis/src/redis-cli shutdown
./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown
fi
./src/common/thirdparty/redis/src/redis-cli shutdown
./src/common/thirdparty/redis/src/redis-cli -p 6380 shutdown