Partial RocksDB React Native AsyncStorage module

// Copyright 2004-present Facebook. All Rights Reserved.

#import "RKAsyncRocksDBStorage.h"

#include

#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::vectorrocksdb::Slice sliceKeys; sliceKeys.reserve(keys.count); for (NSString *key in keys) { sliceKeys.push_back(SliceFromString(key)); } std::vectorstd::string values; std::vectorrocksdb::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