jormun-db/rocksdb/rocksdb.odin

package rocksdb

import "core:c"
import "core:fmt"

foreign import rocksdb "system:rocksdb"

// In order to use RocksDB's WAL replication helpers, we need to import the C++ library so we use this shim
//foreign import rocksdb_shim "system:jormun_rocksdb_shim"  // I know we'll use in future but because we're not right now, compiler is complaining

// RocksDB C API types
RocksDB_T :: distinct rawptr
RocksDB_Options :: distinct rawptr
RocksDB_WriteOptions :: distinct rawptr
RocksDB_ReadOptions :: distinct rawptr
RocksDB_WriteBatch :: distinct rawptr
RocksDB_Iterator :: distinct rawptr
RocksDB_FlushOptions :: distinct rawptr

// Error type
Error :: enum {
	None,
	OpenFailed,
	WriteFailed,
	ReadFailed,
	DeleteFailed,
	InvalidArgument,
	Corruption,
	NotFound,
	IOError,
	Unknown,
}

// Database handle with options
DB :: struct {
	handle:        RocksDB_T,
	options:       RocksDB_Options,
	write_options: RocksDB_WriteOptions,
	read_options:  RocksDB_ReadOptions,
}

// Foreign C functions
@(default_calling_convention = "c")
foreign rocksdb {
	// Database operations
	rocksdb_open :: proc(options: RocksDB_Options, path: cstring, errptr: ^cstring) -> RocksDB_T ---
	rocksdb_close :: proc(db: RocksDB_T) ---

	// Options
	rocksdb_options_create :: proc() -> RocksDB_Options ---
	rocksdb_options_destroy :: proc(options: RocksDB_Options) ---
	rocksdb_options_set_create_if_missing :: proc(options: RocksDB_Options, val: c.uchar) ---
	rocksdb_options_increase_parallelism :: proc(options: RocksDB_Options, total_threads: c.int) ---
	rocksdb_options_optimize_level_style_compaction :: proc(options: RocksDB_Options, memtable_memory_budget: c.uint64_t) ---
	rocksdb_options_set_compression :: proc(options: RocksDB_Options, compression: c.int) ---

	// Write options
	rocksdb_writeoptions_create :: proc() -> RocksDB_WriteOptions ---
	rocksdb_writeoptions_destroy :: proc(options: RocksDB_WriteOptions) ---

	// Read options
	rocksdb_readoptions_create :: proc() -> RocksDB_ReadOptions ---
	rocksdb_readoptions_destroy :: proc(options: RocksDB_ReadOptions) ---

	// Put/Get/Delete
	rocksdb_put :: proc(db: RocksDB_T, options: RocksDB_WriteOptions, key: [^]byte, keylen: c.size_t, val: [^]byte, vallen: c.size_t, errptr: ^cstring) ---
	rocksdb_get :: proc(db: RocksDB_T, options: RocksDB_ReadOptions, key: [^]byte, keylen: c.size_t, vallen: ^c.size_t, errptr: ^cstring) -> [^]byte ---
	rocksdb_delete :: proc(db: RocksDB_T, options: RocksDB_WriteOptions, key: [^]byte, keylen: c.size_t, errptr: ^cstring) ---

	// Flush
	rocksdb_flushoptions_create :: proc() -> RocksDB_FlushOptions ---
	rocksdb_flushoptions_destroy :: proc(options: RocksDB_FlushOptions) ---
	rocksdb_flush :: proc(db: RocksDB_T, options: RocksDB_FlushOptions, errptr: ^cstring) ---

	// Write batch
	rocksdb_writebatch_create :: proc() -> RocksDB_WriteBatch ---
	rocksdb_writebatch_destroy :: proc(batch: RocksDB_WriteBatch) ---
	rocksdb_writebatch_put :: proc(batch: RocksDB_WriteBatch, key: [^]byte, keylen: c.size_t, val: [^]byte, vallen: c.size_t) ---
	rocksdb_writebatch_delete :: proc(batch: RocksDB_WriteBatch, key: [^]byte, keylen: c.size_t) ---
	rocksdb_writebatch_clear :: proc(batch: RocksDB_WriteBatch) ---
	rocksdb_write :: proc(db: RocksDB_T, options: RocksDB_WriteOptions, batch: RocksDB_WriteBatch, errptr: ^cstring) ---

	// Iterator
	rocksdb_create_iterator :: proc(db: RocksDB_T, options: RocksDB_ReadOptions) -> RocksDB_Iterator ---
	rocksdb_iter_destroy :: proc(iter: RocksDB_Iterator) ---
	rocksdb_iter_seek_to_first :: proc(iter: RocksDB_Iterator) ---
	rocksdb_iter_seek_to_last :: proc(iter: RocksDB_Iterator) ---
	rocksdb_iter_seek :: proc(iter: RocksDB_Iterator, key: [^]byte, keylen: c.size_t) ---
	rocksdb_iter_seek_for_prev :: proc(iter: RocksDB_Iterator, key: [^]byte, keylen: c.size_t) ---
	rocksdb_iter_valid :: proc(iter: RocksDB_Iterator) -> c.uchar ---
	rocksdb_iter_next :: proc(iter: RocksDB_Iterator) ---
	rocksdb_iter_prev :: proc(iter: RocksDB_Iterator) ---
	rocksdb_iter_key :: proc(iter: RocksDB_Iterator, klen: ^c.size_t) -> [^]byte ---
	rocksdb_iter_value :: proc(iter: RocksDB_Iterator, vlen: ^c.size_t) -> [^]byte ---

	// Memory management
	rocksdb_free :: proc(ptr: rawptr) ---
}

// Compression types
ROCKSDB_NO_COMPRESSION :: 0
ROCKSDB_SNAPPY_COMPRESSION :: 1
ROCKSDB_ZLIB_COMPRESSION :: 2
ROCKSDB_BZIP2_COMPRESSION :: 3
ROCKSDB_LZ4_COMPRESSION :: 4
ROCKSDB_LZ4HC_COMPRESSION :: 5
ROCKSDB_ZSTD_COMPRESSION :: 7

// Open a database
db_open :: proc(path: string, create_if_missing := true) -> (DB, Error) {
	options := rocksdb_options_create()
	if options == nil {
		return {}, .Unknown
	}

	// Set create if missing
	rocksdb_options_set_create_if_missing(options, create_if_missing ? 1 : 0)

	// Performance optimizations
	rocksdb_options_increase_parallelism(options, 4)
	rocksdb_options_optimize_level_style_compaction(options, 512 * 1024 * 1024)
	rocksdb_options_set_compression(options, ROCKSDB_LZ4_COMPRESSION)

	// Create write and read options
	write_options := rocksdb_writeoptions_create()
	if write_options == nil {
		rocksdb_options_destroy(options)
		return {}, .Unknown
	}

	read_options := rocksdb_readoptions_create()
	if read_options == nil {
		rocksdb_writeoptions_destroy(write_options)
		rocksdb_options_destroy(options)
		return {}, .Unknown
	}

	// Open database
	err: cstring
	path_cstr := fmt.ctprintf("%s", path)
	handle := rocksdb_open(options, path_cstr, &err)
	if err != nil {
		defer rocksdb_free(rawptr(err)) // Cast it here and now so we don't deal with issues from FFI down the line
		rocksdb_readoptions_destroy(read_options)
		rocksdb_writeoptions_destroy(write_options)
		rocksdb_options_destroy(options)
		return {}, .OpenFailed
	}

	return DB{
		handle = handle,
		options = options,
		write_options = write_options,
		read_options = read_options,
	}, .None
}

// Close database
db_close :: proc(db: ^DB) {
	rocksdb_readoptions_destroy(db.read_options)
	rocksdb_writeoptions_destroy(db.write_options)
	rocksdb_close(db.handle)
	rocksdb_options_destroy(db.options)
}

// Put key-value pair
db_put :: proc(db: ^DB, key: []byte, value: []byte) -> Error {
	err: cstring
	rocksdb_put(
		db.handle,
		db.write_options,
		raw_data(key),
		c.size_t(len(key)),
		raw_data(value),
		c.size_t(len(value)),
		&err,
	)
	if err != nil {
		defer rocksdb_free(rawptr(err)) // Cast it here and now so we don't deal with issues from FFI down the line
		return .WriteFailed
	}
	return .None
}

// Get value by key (returns owned slice - caller must free)
db_get :: proc(db: ^DB, key: []byte) -> (value: []byte, err: Error) {
	errptr: cstring
	value_len: c.size_t

	value_ptr := rocksdb_get(
		db.handle,
		db.read_options,
		raw_data(key),
		c.size_t(len(key)),
		&value_len,
		&errptr,
	)

	if errptr != nil {
		defer rocksdb_free(rawptr(errptr)) // Cast it here and now so we don't deal with issues from FFI down the line
		return nil, .ReadFailed
	}

	if value_ptr == nil {
		return nil, .NotFound
	}

	// Copy the data and free RocksDB's buffer
	result := make([]byte, value_len, context.allocator)
	copy(result, value_ptr[:value_len])
	rocksdb_free(rawptr(value_ptr)) // Cast it here and now so we don't deal with issues from FFI down the line

	return result, .None
}

// Delete key
db_delete :: proc(db: ^DB, key: []byte) -> Error {
	err: cstring
	rocksdb_delete(
		db.handle,
		db.write_options,
		raw_data(key),
		c.size_t(len(key)),
		&err,
	)
	if err != nil {
		defer rocksdb_free(rawptr(err)) // Cast it here and now so we don't deal with issues from FFI down the line
		return .DeleteFailed
	}
	return .None
}

// Flush database
db_flush :: proc(db: ^DB) -> Error {
	flush_opts := rocksdb_flushoptions_create()
	if flush_opts == nil {
		return .Unknown
	}
	defer rocksdb_flushoptions_destroy(flush_opts)

	err: cstring
	rocksdb_flush(db.handle, flush_opts, &err)
	if err != nil {
		defer rocksdb_free(rawptr(err)) // Cast it here and now so we don't deal with issues from FFI down the line
		return .IOError
	}
	return .None
}

// Write batch
WriteBatch :: struct {
	handle: RocksDB_WriteBatch,
}

// Create write batch
batch_create :: proc() -> (WriteBatch, Error) {
	handle := rocksdb_writebatch_create()
	if handle == nil {
		return {}, .Unknown
	}
	return WriteBatch{handle = handle}, .None
}

// Destroy write batch
batch_destroy :: proc(batch: ^WriteBatch) {
	rocksdb_writebatch_destroy(batch.handle)
}

// Add put operation to batch
batch_put :: proc(batch: ^WriteBatch, key: []byte, value: []byte) {
	rocksdb_writebatch_put(
		batch.handle,
		raw_data(key),
		c.size_t(len(key)),
		raw_data(value),
		c.size_t(len(value)),
	)
}

// Add delete operation to batch
batch_delete :: proc(batch: ^WriteBatch, key: []byte) {
	rocksdb_writebatch_delete(
		batch.handle,
		raw_data(key),
		c.size_t(len(key)),
	)
}

// Clear batch
batch_clear :: proc(batch: ^WriteBatch) {
	rocksdb_writebatch_clear(batch.handle)
}

// Write batch to database
batch_write :: proc(db: ^DB, batch: ^WriteBatch) -> Error {
	err: cstring
	rocksdb_write(db.handle, db.write_options, batch.handle, &err)
	if err != nil {
		defer rocksdb_free(rawptr(err)) // Cast it here and now so we don't deal with issues from FFI down the line
		return .WriteFailed
	}
	return .None
}

// Iterator
Iterator :: struct {
	handle: RocksDB_Iterator,
}

// Create iterator
iter_create :: proc(db: ^DB) -> (Iterator, Error) {
	handle := rocksdb_create_iterator(db.handle, db.read_options)
	if handle == nil {
		return {}, .Unknown
	}
	return Iterator{handle = handle}, .None
}

// Destroy iterator
iter_destroy :: proc(iter: ^Iterator) {
	rocksdb_iter_destroy(iter.handle)
}

// Seek to first
iter_seek_to_first :: proc(iter: ^Iterator) {
	rocksdb_iter_seek_to_first(iter.handle)
}

// Seek to last
iter_seek_to_last :: proc(iter: ^Iterator) {
	rocksdb_iter_seek_to_last(iter.handle)
}

// Seek to key
iter_seek :: proc(iter: ^Iterator, target: []byte) {
	rocksdb_iter_seek(iter.handle, raw_data(target), c.size_t(len(target)))
}

// Check if iterator is valid
iter_valid :: proc(iter: ^Iterator) -> bool {
	return rocksdb_iter_valid(iter.handle) != 0
}

// Move to next
iter_next :: proc(iter: ^Iterator) {
	rocksdb_iter_next(iter.handle)
}

// Move to previous
iter_prev :: proc(iter: ^Iterator) {
	rocksdb_iter_prev(iter.handle)
}

// Get current key (returns borrowed slice)
iter_key :: proc(iter: ^Iterator) -> []byte {
	klen: c.size_t
	key_ptr := rocksdb_iter_key(iter.handle, &klen)
	if key_ptr == nil {
		return nil
	}
	return key_ptr[:klen]
}

// Get current value (returns borrowed slice)
iter_value :: proc(iter: ^Iterator) -> []byte {
	vlen: c.size_t
	value_ptr := rocksdb_iter_value(iter.handle, &vlen)
	if value_ptr == nil {
		return nil
	}
	return value_ptr[:vlen]
}