Skip to content

Instantly share code, notes, and snippets.

@sahrens
Created October 14, 2015 19:57
Show Gist options
  • Save sahrens/ae3ad0889c608ecd51aa to your computer and use it in GitHub Desktop.
Save sahrens/ae3ad0889c608ecd51aa to your computer and use it in GitHub Desktop.
Partial RocksDB React Native AsyncStorage module
// Copyright 2004-present Facebook. All Rights Reserved.
#import "RKAsyncRocksDBStorage.h"
#include <string>
#import <Foundation/Foundation.h>
#import <FBReactKit/RCTConvert.h>
#import <FBReactKit/RCTLog.h>
#import <FBReactKit/RCTUtils.h>
#define ROCKSDB_LITE 1
#include <rocksdb/db.h>
#include <rocksdb/merge_operator.h>
#include <rocksdb/options.h>
#include <rocksdb/slice.h>
#include <rocksdb/status.h>
#import <SSZipArchive/SSZipArchive.h>
static NSString *const RKAsyncRocksDBStorageDirectory = @"RKAsyncRocksDBStorage";
namespace {
rocksdb::Slice SliceFromString(NSString *string)
{
return rocksdb::Slice((const char *)[string UTF8String], [string lengthOfBytesUsingEncoding:NSUTF8StringEncoding]);
}
void deepMergeInto(NSMutableDictionary *output, NSDictionary *input) {
for (NSString *key in input) {
id inputValue = input[key];
if ([inputValue isKindOfClass:[NSDictionary class]]) {
NSMutableDictionary *nestedOutput;
id outputValue = output[key];
if ([outputValue isKindOfClass:[NSMutableDictionary class]]) {
nestedOutput = outputValue;
} else {
if ([outputValue isKindOfClass:[NSDictionary class]]) {
nestedOutput = [outputValue mutableCopy];
} else {
output[key] = [inputValue copy];
}
}
if (nestedOutput) {
deepMergeInto(nestedOutput, inputValue);
output[key] = nestedOutput;
}
} else {
output[key] = inputValue;
}
}
}
class JSONMergeOperator : public rocksdb::AssociativeMergeOperator {
public:
virtual bool Merge(const rocksdb::Slice &key,
const rocksdb::Slice *existingValue,
const rocksdb::Slice &newValue,
std::string *mergedValue,
rocksdb::Logger *logger) const override {
NSError *error;
NSMutableDictionary *existingDict;
if (existingValue) {
NSString *existingString = [NSString stringWithCString:existingValue->data() length:existingValue->size()];
existingDict = RCTJSONParseMutable(existingString, &error);
if (error) {
RCTLogError(@"Parse error in RKAsyncRocksDBStorage merge operation. Error:\n%@\nString:\n%@", error, existingString);
return false;
}
} else {
// Nothing to merge, just assign the string without even parsing.
mergedValue->assign(newValue.data(), newValue.size());
return true;
}
NSString *newString = [NSString stringWithCString:newValue.data() length:newValue.size()];
NSMutableDictionary *newDict = RCTJSONParse(newString, &error);
deepMergeInto(existingDict, newDict);
NSString *mergedNSString = RCTJSONStringify(existingDict, &error);
mergedValue->assign([mergedNSString UTF8String], [mergedNSString lengthOfBytesUsingEncoding:NSUTF8StringEncoding]);
return true;
}
virtual const char *Name() const override {
return "JSONMergeOperator";
}
};
} // namespace
@implementation RKAsyncRocksDBStorage
{
rocksdb::DB *_db;
}
@synthesize methodQueue = _methodQueue;
namespace {
RCT_EXPORT_MODULE()
- (BOOL)ensureDirectorySetup:(NSError **)error
{
if (_db) {
return YES;
}
RCTAssert(error != nil, @"Must provide error pointer.");
rocksdb::Status status = rocksdb::DB::Open(options, [userPath UTF8String], &_db);
if (!status.ok() || !_db) {
RCTLogError(@"Failed to open db at path %@.\n\nRocksDB Status: %s.\n\nNSError: %@", userPath, status.ToString().c_str(), *error);
*error = [NSError errorWithDomain:@"rocksdb" code:100 userInfo:@{NSLocalizedDescriptionKey:@"Failed to open db"}];
return NO;
}
return YES;
}
RCT_EXPORT_METHOD(multiGet:(NSStringArray *)keys
callback:(RCTResponseSenderBlock)callback)
{
NSDictionary *errorOut;
NSError *error;
NSMutableArray *result = [[NSMutableArray alloc] initWithCapacity:keys.count];
BOOL success = [self ensureDirectorySetup:&error];
if (!success || error) {
errorOut = RCTMakeError(@"Failed to setup directory", nil, nil);
} else {
std::vector<rocksdb::Slice> sliceKeys;
sliceKeys.reserve(keys.count);
for (NSString *key in keys) {
sliceKeys.push_back(SliceFromString(key));
}
std::vector<std::string> values;
std::vector<rocksdb::Status> statuses = _db->MultiGet(rocksdb::ReadOptions(), sliceKeys, &values);
RCTAssert(values.size() == keys.count, @"Key and value arrays should be equal size");
for (size_t ii = 0; ii < values.size(); ii++) {
id value = [NSNull null];
auto status = statuses[ii];
if (!status.IsNotFound()) {
if (!status.ok()) {
errorOut = RCTMakeError(@"RKAsyncRocksDB failed getting key: ", keys[ii], keys[ii]);
} else {
value = [NSString stringWithUTF8String:values[ii].c_str()];
}
}
[result addObject:@[keys[ii], value]];
}
}
if (callback) {
callback(@[errorOut ? @[errorOut] : [NSNull null], result]);
}
}
// kvPairs is a list of key-value pairs, e.g. @[@[key1, val1], @[key2, val2], ...]
// TODO: write custom RCTConvert category method for kvPairs
RCT_EXPORT_METHOD(multiSet:(NSArray *)kvPairs
callback:(RCTResponseSenderBlock)callback)
{
auto updates = rocksdb::WriteBatch();
for (NSArray *kvPair in kvPairs) {
NSStringArray *pair = [RCTConvert NSStringArray:kvPair];
if (pair.count == 2) {
updates.Put(SliceFromString(kvPair[0]), SliceFromString(kvPair[1]));
} else {
if (callback) {
callback(@[@[RCTMakeAndLogError(@"Input must be an array of [key, value] arrays, got: ", kvPair, nil)]]);
}
return;
}
}
[self _performWriteBatch:&updates callback:callback];
}
RCT_EXPORT_METHOD(multiMerge:(NSArray *)kvPairs
callback:(RCTResponseSenderBlock)callback)
{
auto updates = rocksdb::WriteBatch();
for (NSArray *kvPair in kvPairs) {
NSStringArray *pair = [RCTConvert NSStringArray:kvPair];
if (pair.count == 2) {
updates.Merge(SliceFromString(pair[0]), SliceFromString(pair[1]));
} else {
if (callback) {
callback(@[@[RCTMakeAndLogError(@"Input must be an array of [key, value] arrays, got: ", kvPair, nil)]]);
}
return;
}
}
[self _performWriteBatch:&updates callback:callback];
}
RCT_EXPORT_METHOD(multiRemove:(NSArray *)keys
callback:(RCTResponseSenderBlock)callback)
{
auto updates = rocksdb::WriteBatch();
for (NSString *key in keys) {
updates.Delete(SliceFromString(key));
}
[self _performWriteBatch:&updates callback:callback];
}
// TODO (#5906496): There's a lot of duplication in the error handling code here - can we refactor this?
- (void)_performWriteBatch:(rocksdb::WriteBatch *)updates callback:(RCTResponseSenderBlock)callback
{
NSDictionary *errorOut;
NSError *error;
BOOL success = [self ensureDirectorySetup:&error];
if (!success || error) {
errorOut = RCTMakeError(@"Failed to setup storage", nil, nil);
} else {
rocksdb::Status status = _db->Write(rocksdb::WriteOptions(), updates);
if (!status.ok()) {
errorOut = RCTMakeError(@"Failed to write to RocksDB database.", nil, nil);
}
}
if (callback) {
callback(@[errorOut ? @[errorOut] : [NSNull null]]);
}
}
RCT_EXPORT_METHOD(clear:(RCTResponseSenderBlock)callback)
{
[self _nullOutDB];
NSDictionary *errorOut;
NSError *error;
NSURL *userDirectory = getOrCreateRocksDBPath(&error);
if (!userDirectory) {
errorOut = RCTMakeError(@"Failed to setup storage", nil, nil);
} else {
rocksdb::Status status = rocksdb::DestroyDB([[userDirectory path] UTF8String], rocksdb::Options());
if (!status.ok()) {
errorOut = RCTMakeError(@"RocksDB:clear failed to destroy db at path ", [userDirectory path], nil);
}
}
if (callback) {
callback(@[errorOut ?: [NSNull null]]);
}
}
RCT_EXPORT_METHOD(getAllKeys:(RCTResponseSenderBlock)callback)
{
NSError *error;
NSMutableArray *allKeys = [NSMutableArray new];
NSDictionary *errorOut;
BOOL success = [self ensureDirectorySetup:&error];
if (!success || error) {
errorOut = RCTMakeError(@"Failed to setup storage", nil, nil);
} else {
rocksdb::Iterator *it = _db->NewIterator(rocksdb::ReadOptions());
for (it->SeekToFirst(); it->Valid(); it->Next()) {
[allKeys addObject:[NSString stringWithCString:it->key().data() length:it->key().size()]];
}
}
if (callback) {
callback(@[errorOut ?: [NSNull null], allKeys]);
}
}
@end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment