Partial RocksDB React Native AsyncStorage module

gistfile1.txt

// 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