Disk: ReadDriver: Add unified cache

Previously, we had to copy an entire block out of the
old cache every time we wanted to read even a single
byte from it.

This ended up being a fairly significant performance
issue, in addition to the fact that the caching code
was duplicated.

disk/diskreaddriver.cpp

@@ -9,14 +9,25 @@ optional<uint64_t> ReadDriver::readLogicalDisk(Communication& com, device_eventh
 	if(amount == 0)
 		return 0;
 
-	optional<uint64_t> ret;
+	pos += vsaOffset;
+
+	// First read from the cache
+	optional<uint64_t> ret = readFromCache(pos, into, amount);
+	if(ret == amount) // Full cache hit, we're done
+		return ret;
+
+	const uint64_t totalAmount = amount;
+	if(ret.has_value()) { // Partial cache hit
+		pos += *ret;
+		into += *ret;
+		amount -= *ret;
+	}
 
 	// Read into here if we can't read directly into the user buffer
 	// That would be the case either if we don't want some at the
 	// beginning or end of the block.
 	std::vector<uint8_t> alignedReadBuffer;
 
-	pos += vsaOffset;
 	const uint32_t idealBlockSize = getBlockSizeBounds().second;
 	const uint64_t startBlock = pos / idealBlockSize;
 	const uint32_t posWithinFirstBlock = static_cast<uint32_t>(pos % idealBlockSize);
@@ -36,7 +47,7 @@ optional<uint64_t> ReadDriver::readLogicalDisk(Communication& com, device_eventh
 
 		const uint32_t posWithinCurrentBlock = (blocksProcessed ? 0 : posWithinFirstBlock);
 		uint32_t curAmt = idealBlockSize - posWithinCurrentBlock;
-		const auto amountLeft = amount - ret.value_or(0);
+		const auto amountLeft = totalAmount - ret.value_or(0);
 		if(curAmt > amountLeft)
 			curAmt = static_cast<uint32_t>(amountLeft);
 
@@ -65,7 +76,41 @@ optional<uint64_t> ReadDriver::readLogicalDisk(Communication& com, device_eventh
 			ret.emplace();
 		*ret += std::min<uint64_t>(*readAmount, curAmt);
 		blocksProcessed++;
+
+		if(blocksProcessed == blocks) {
+			// Last block, add to the cache
+			if(useAlignedReadBuffer) {
+				cache = std::move(alignedReadBuffer);
+			} else {
+				if(cache.size() != idealBlockSize)
+					cache.resize(idealBlockSize);
+				memcpy(cache.data(), into + intoOffset, idealBlockSize);
+			}
+			cachePos = currentBlock * idealBlockSize;
+			cachedAt = std::chrono::steady_clock::now();
+		}
 	}
 
 	return ret;
 }
+
+void ReadDriver::invalidateCache(uint64_t pos, uint64_t amount) {
+	if(pos <= cachePos + cache.size() && pos + amount >= cachePos)
+		cache.clear();
+}
+
+optional<uint64_t> ReadDriver::readFromCache(uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds staleAfter) {
+	if(cache.empty())
+		return nullopt; // Nothing in the cache
+
+	if(cachedAt + staleAfter < std::chrono::steady_clock::now())
+		return nullopt; // Cache is stale
+	
+	if(pos > cachePos + cache.size() || pos < cachePos)
+		return nullopt; // Cache miss
+
+	const auto cacheOffset = pos - cachePos;
+	const auto copyAmount = std::min<uint64_t>(cache.size() - cacheOffset, amount);
+	memcpy(into, cache.data() + cacheOffset, static_cast<size_t>(copyAmount));
+	return copyAmount;
+}

disk/diskwritedriver.cpp

@@ -75,6 +75,7 @@ optional<uint64_t> WriteDriver::writeLogicalDisk(Communication& com, device_even
 		if(bytesTransferred == RetryAtomic) {
 			// The user may want to log these events in order to see how many atomic misses they are getting
 			report(APIEvent::Type::AtomicOperationRetried, APIEvent::Severity::EventInfo);
+			readDriver.invalidateCache(currentBlock * idealBlockSize, idealBlockSize);
 			continue;
 		}
 
@@ -93,5 +94,9 @@ optional<uint64_t> WriteDriver::writeLogicalDisk(Communication& com, device_even
 		blocksProcessed++;
 	}
 
+	// No matter how much succeeded, to be safe, we'll invalidate anything
+	// we may have even tried to write, since it may have succeeded without
+	// notifying, etc.
+	readDriver.invalidateCache(pos, amount);
 	return ret;
 }

disk/extextractordiskreaddriver.cpp

@@ -37,7 +37,6 @@ optional<uint64_t> ExtExtractorDiskReadDriver::attemptReadLogicalDiskAligned(Com
 	uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout) {
 	static std::shared_ptr<MessageFilter> NeoMemorySDRead = std::make_shared<MessageFilter>(Network::NetID::NeoMemorySDRead);
 
-	if(cachePos != pos || std::chrono::steady_clock::now() > cachedAt + CacheTime) {
 	uint64_t sector = pos / SectorSize;
 
 	uint64_t largeSectorCount = amount / SectorSize;
@@ -45,7 +44,6 @@ optional<uint64_t> ExtExtractorDiskReadDriver::attemptReadLogicalDiskAligned(Com
 	if (largeSectorCount != uint64_t(sectorCount))
 		return nullopt;
 
-		// The cache does not have this data, go get it
 	std::mutex m;
 	std::condition_variable cv;
 	uint16_t receiving = 0; // How much are we about to get before another header or completion
@@ -146,7 +144,7 @@ optional<uint64_t> ExtExtractorDiskReadDriver::attemptReadLogicalDiskAligned(Com
 			#ifdef ICSNEO_EXTENDED_EXTRACTOR_DEBUG_PRINTS
 			std::cout << "With " << int(left) << " bytes " << int(offset) << std::endl;
 			#endif
-				memcpy(cache.data() + received, data.data() + offset, count);
+			memcpy(into + received, data.data() + offset, count);
 			received += count;
 			receivedCurrent += count;
 			offset += count;
@@ -204,10 +202,5 @@ optional<uint64_t> ExtExtractorDiskReadDriver::attemptReadLogicalDiskAligned(Com
 	if(hitTimeout || error)
 		return nullopt;
 
-		cachedAt = std::chrono::steady_clock::now();
-		cachePos = pos;
-	}
-
-	memcpy(into, cache.data(), size_t(amount));
 	return amount;
 }

disk/neomemorydiskdriver.cpp

@@ -15,8 +15,6 @@ optional<uint64_t> NeoMemoryDiskDriver::readLogicalDiskAligned(Communication& co
 	if(amount != SectorSize)
 		return nullopt;
 
-	if(cachePos != pos || std::chrono::steady_clock::now() > cachedAt + CacheTime) {
-		// The cache does not have this data, go get it
 	const uint64_t currentSector = pos / SectorSize;
 	auto msg = com.waitForMessageSync([&currentSector, &com] {
 		return com.sendCommand(Command::NeoReadMemory, {
@@ -41,12 +39,7 @@ optional<uint64_t> NeoMemoryDiskDriver::readLogicalDiskAligned(Communication& co
 		return nullopt;
 	}
 
-		memcpy(cache.data(), sdmsg->data.data(), SectorSize);
+	memcpy(into, sdmsg->data.data(), SectorSize);
-		cachedAt = std::chrono::steady_clock::now();
-		cachePos = pos;
-	}
-
-	memcpy(into, cache.data(), SectorSize);
 	return SectorSize;
 }
 
@@ -61,10 +54,6 @@ optional<uint64_t> NeoMemoryDiskDriver::writeLogicalDiskAligned(Communication& c
 	if(amount != SectorSize)
 		return nullopt;
 
-	// Clear the cache if we're writing to the cached sector
-	if(pos == cachePos)
-		cachedAt = std::chrono::time_point<std::chrono::steady_clock>();
-
 	// Requesting an atomic operation, but neoMemory does not support it
 	// Continue on anyway but warn the caller
 	if(atomicBuf != nullptr)

disk/plasiondiskreaddriver.cpp

@@ -19,13 +19,11 @@ optional<uint64_t> PlasionDiskReadDriver::readLogicalDiskAligned(Communication&
 	if(pos % getBlockSizeBounds().first != 0)
 		return nullopt;
 
-	if(cachePos != pos || std::chrono::steady_clock::now() > cachedAt + CacheTime) {
 	uint64_t largeSector = pos / SectorSize;
 	uint32_t sector = uint32_t(largeSector);
 	if (largeSector != uint64_t(sector))
 		return nullopt;
 
-		// The cache does not have this data, go get it
 	std::mutex m;
 	std::condition_variable cv;
 	uint32_t copied = 0;
@@ -35,17 +33,14 @@ optional<uint64_t> PlasionDiskReadDriver::readLogicalDiskAligned(Communication&
 		std::unique_lock<std::mutex> lk(m);
 
 		const auto sdmsg = std::dynamic_pointer_cast<NeoReadMemorySDMessage>(msg);
-			if(!sdmsg || cache.size() < copied + sdmsg->data.size()) {
+		if(!sdmsg || amount < copied + sdmsg->data.size()) {
 			error = true;
 			lk.unlock();
 			cv.notify_all();
 			return;
 		}
 
-			// Invalidate the cache here in case we fail half-way through
+		memcpy(into + copied, sdmsg->data.data(), sdmsg->data.size());
-			cachedAt = std::chrono::steady_clock::time_point();
-
-			memcpy(cache.data() + copied, sdmsg->data.data(), sdmsg->data.size());
 		copied += uint32_t(sdmsg->data.size());
 		if(copied == amount) {
 			lk.unlock();
@@ -69,10 +64,5 @@ optional<uint64_t> PlasionDiskReadDriver::readLogicalDiskAligned(Communication&
 	if(hitTimeout)
 		return nullopt;
 
-		cachedAt = std::chrono::steady_clock::now();
-		cachePos = pos;
-	}
-
-	memcpy(into, cache.data(), size_t(amount));
 	return amount;
 }

include/icsneo/disk/diskreaddriver.h

@@ -22,6 +22,9 @@ public:
 	virtual optional<uint64_t> readLogicalDisk(Communication& com, device_eventhandler_t report,
 		uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout = DefaultTimeout);
 
+	void invalidateCache(uint64_t pos = 0,
+		uint64_t amount = std::numeric_limits<uint32_t>::max() /* large value, but avoid overflow */);
+
 protected:
 	/**
 	 * Perform a read which the driver can do in one shot.
@@ -31,6 +34,15 @@ protected:
 	 */
 	virtual optional<uint64_t> readLogicalDiskAligned(Communication& com, device_eventhandler_t report,
 		uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds timeout) = 0;
+
+private:
+	std::vector<uint8_t> cache;
+	uint64_t cachePos = 0;
+	std::chrono::time_point<std::chrono::steady_clock> cachedAt;
+
+	static constexpr const std::chrono::milliseconds CacheTime = std::chrono::milliseconds(1000);
+
+	optional<uint64_t> readFromCache(uint64_t pos, uint8_t* into, uint64_t amount, std::chrono::milliseconds staleAfter = CacheTime);
 };
 
 } // namespace Disk

include/icsneo/disk/extextractordiskreaddriver.h

@@ -24,13 +24,8 @@ public:
 
 private:
 	static constexpr const uint32_t MaxSize = Disk::SectorSize * 512;
-	static constexpr const std::chrono::seconds CacheTime = std::chrono::seconds(1);
 	static constexpr const uint8_t HeaderLength = 7;
 
-	std::array<uint8_t, MaxSize> cache;
-	uint64_t cachePos = 0;
-	std::chrono::time_point<std::chrono::steady_clock> cachedAt;
-
 	Access getPossibleAccess() const override { return Access::EntireCard; }
 
 	optional<uint64_t> readLogicalDiskAligned(Communication& com, device_eventhandler_t report,

include/icsneo/disk/neomemorydiskdriver.h

@@ -27,11 +27,6 @@ public:
 
 private:
 	static constexpr const uint8_t MemoryTypeSD = 0x01; // Logical Disk
-	static constexpr const std::chrono::seconds CacheTime = std::chrono::seconds(1);
-
-	std::array<uint8_t, SectorSize> cache;
-	uint64_t cachePos = 0;
-	std::chrono::time_point<std::chrono::steady_clock> cachedAt;
 
 	Access getPossibleAccess() const override { return Access::VSA; }
 

include/icsneo/disk/plasiondiskreaddriver.h

@@ -24,11 +24,6 @@ public:
 
 private:
 	static constexpr const uint32_t MaxSize = 65024;
-	static constexpr const std::chrono::seconds CacheTime = std::chrono::seconds(1);
-
-	std::array<uint8_t, MaxSize> cache;
-	uint64_t cachePos = 0;
-	std::chrono::time_point<std::chrono::steady_clock> cachedAt;
 
 	Access getPossibleAccess() const override { return Access::EntireCard; }
 

test/diskdriverreadtest.cpp

@@ -62,3 +62,55 @@ TEST_F(DiskDriverTest, ReadBadStartingPos) {
 	EXPECT_FALSE(amountRead.has_value());
 	EXPECT_EQ(driver->readCalls, 1u); // One to check EOF
 }
+
+TEST_F(DiskDriverTest, ReadCache) {
+	std::array<uint8_t, 128> buf;
+	buf.fill(0u);
+	auto amountRead = readLogicalDisk(1, buf.data(), buf.size());
+	EXPECT_TRUE(amountRead.has_value());
+	EXPECT_EQ(amountRead, buf.size());
+	EXPECT_EQ(buf[0], TEST_STRING[1]);
+	EXPECT_EQ(buf[110], 111u);
+	EXPECT_EQ(driver->readCalls, 1u);
+
+	// Subsequent reads (within the same second) should hit the cache
+	amountRead = readLogicalDisk(1, buf.data(), buf.size());
+	EXPECT_EQ(driver->readCalls, 1u);
+
+	// The underlying data can be changed
+	driver->mockDisk[1] = 'J';
+
+	// But the same data should be returned from the cache
+	amountRead = readLogicalDisk(1, buf.data(), buf.size());
+	EXPECT_TRUE(amountRead.has_value());
+	EXPECT_EQ(amountRead, buf.size());
+	EXPECT_EQ(buf[0], TEST_STRING[1]);
+	EXPECT_EQ(buf[110], 111u);
+	EXPECT_EQ(driver->readCalls, 1u);
+
+	driver->invalidateCache(0, 0xfffff);
+
+	// After invalidating the cache (or waiting for it to expire), the underlying data will be read
+	amountRead = readLogicalDisk(1, buf.data(), buf.size());
+	EXPECT_TRUE(amountRead.has_value());
+	EXPECT_EQ(amountRead, buf.size());
+	EXPECT_EQ(buf[0], 'J');
+	EXPECT_EQ(buf[110], 111u);
+	EXPECT_EQ(driver->readCalls, 2u);
+}
+
+TEST_F(DiskDriverTest, ReadCacheLong) {
+	std::array<uint8_t, 500> buf;
+	buf.fill(0u);
+	auto amountRead = readLogicalDisk(300, buf.data(), buf.size());
+	EXPECT_TRUE(amountRead.has_value());
+	EXPECT_EQ(amountRead, buf.size());
+	EXPECT_EQ(buf[0], 300 & 0xFF);
+	EXPECT_EQ(buf[110], 410 & 0xFF);
+	EXPECT_EQ(driver->readCalls, 3u);
+
+	// Re-read the end, it will be in the cache
+	amountRead = readLogicalDisk(780, buf.data() + 480, buf.size() - 480);
+	EXPECT_EQ(buf[490], 790 & 0xFF);
+	EXPECT_EQ(driver->readCalls, 3u);
+}