From 4e14ae77db9288a42418ce23bbe5367a741c1fac Mon Sep 17 00:00:00 2001
From: Andrea Cavalli <andrea@cavallium.it>
Date: Wed, 25 Sep 2024 14:58:38 +0200
Subject: [PATCH] Avoid crashes during sst errors

---
 .../java/it/cavallium/rockserver/core/impl/EmbeddedDB.java   | 5 +++++
 .../it/cavallium/rockserver/core/impl/rocksdb/SSTWriter.java | 2 --
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/main/java/it/cavallium/rockserver/core/impl/EmbeddedDB.java b/src/main/java/it/cavallium/rockserver/core/impl/EmbeddedDB.java
index 9cec5d1..87cf27b 100644
--- a/src/main/java/it/cavallium/rockserver/core/impl/EmbeddedDB.java
+++ b/src/main/java/it/cavallium/rockserver/core/impl/EmbeddedDB.java
@@ -440,6 +440,7 @@ public class EmbeddedDB implements RocksDBSyncAPI, Closeable {
 		try {
 			var cf = new CompletableFuture<Void>();
 			batchPublisher.subscribe(new Subscriber<>() {
+				private boolean stopped;
 				private Subscription subscription;
 				private ColumnInstance col;
 				private ArrayList<AutoCloseable> refs;
@@ -476,6 +477,9 @@ public class EmbeddedDB implements RocksDBSyncAPI, Closeable {
 
 				@Override
 				public void onNext(KVBatch kvBatch) {
+					if (stopped) {
+						return;
+					}
 					var keyIt = kvBatch.keys().iterator();
 					var valueIt = kvBatch.values().iterator();
 					try (var arena = Arena.ofConfined()) {
@@ -516,6 +520,7 @@ public class EmbeddedDB implements RocksDBSyncAPI, Closeable {
 				}
 
 				private void doFinally() {
+					stopped = true;
 					for (int i = refs.size() - 1; i >= 0; i--) {
 						try {
 							var c = refs.get(i);
diff --git a/src/main/java/it/cavallium/rockserver/core/impl/rocksdb/SSTWriter.java b/src/main/java/it/cavallium/rockserver/core/impl/rocksdb/SSTWriter.java
index 80eaa65..6f1cee3 100644
--- a/src/main/java/it/cavallium/rockserver/core/impl/rocksdb/SSTWriter.java
+++ b/src/main/java/it/cavallium/rockserver/core/impl/rocksdb/SSTWriter.java
@@ -85,7 +85,6 @@ public record SSTWriter(RocksDB db, it.cavallium.rockserver.core.impl.ColumnInst
 
     public void put(byte[] key, byte[] value) throws RocksDBException {
         try {
-            checkOwningHandle();
             sstFileWriter.put(key, value);
         } catch (org.rocksdb.RocksDBException e) {
             throw RocksDBException.of(RocksDBException.RocksDBErrorType.PUT_UNKNOWN_ERROR, e);
@@ -94,7 +93,6 @@ public record SSTWriter(RocksDB db, it.cavallium.rockserver.core.impl.ColumnInst
 
     public void put(ByteBuffer key, ByteBuffer value) throws RocksDBException {
         try {
-            checkOwningHandle();
             sstFileWriter.put(key, value);
         } catch (org.rocksdb.RocksDBException e) {
             throw RocksDBException.of(RocksDBException.RocksDBErrorType.PUT_UNKNOWN_ERROR, e);