foundationdb/fdbclient/AsyncFileBlobStore.h

/*
 * AsyncFileBlobStore.h
 *
 * This source file is part of the FoundationDB open source project
 *
 * Copyright 2013-2026 Apple Inc. and the FoundationDB project authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <sstream>
#include <time.h>

#include "flow/IAsyncFile.h"
#include "flow/serialize.h"
#include "flow/Net2Packet.h"
#include "flow/IRateControl.h"
#include "fdbclient/IBlobStore.h"
#include "md5/md5.h"
#include "libb64/encode.h"
#include <openssl/sha.h>

template <typename T>
static Future<T> joinErrorGroup(Future<T> f, Promise<Void> p) {
	try {
		co_await (success(f) || p.getFuture());
		co_return f.get();
	} catch (Error& e) {
		if (p.canBeSet())
			p.sendError(e);
		throw;
	}
}
// This class represents a write-only file that lives in a blob store.  It writes using the REST API,
// using multi-part upload and beginning to transfer each part as soon as it is large enough.
// All write operations file operations must be sequential and contiguous.
// Limits on part sizes, upload speed, and concurrent uploads are taken from the IBlobStoreEndpoint being used.
class AsyncFileBlobStoreWrite final : public IAsyncFile, public ReferenceCounted<AsyncFileBlobStoreWrite> {
public:
	void addref() override { ReferenceCounted<AsyncFileBlobStoreWrite>::addref(); }
	void delref() override { ReferenceCounted<AsyncFileBlobStoreWrite>::delref(); }

	StringRef getClassName() override { return "AsyncFileBlobStoreWrite"_sr; }

	struct Part : ReferenceCounted<Part> {
		Part(int n, int minSize, bool useSHA256)
		  : number(n), writer(content.getWriteBuffer(minSize), nullptr, Unversioned()), length(0),
		    use_sha256(useSHA256) {
			etag = std::string();
			if (use_sha256) {
				::SHA256_Init(&content_sha256_buf);
			} else {
				::MD5_Init(&content_md5_buf);
			}
		}
		virtual ~Part() { etag.cancel(); }
		Future<std::string> etag;
		int number;
		UnsentPacketQueue content;
		std::string checksumString; // Contains either MD5 or SHA256 based on use_sha256
		PacketWriter writer;
		int length;
		void write(const uint8_t* buf, int len) {
			writer.serializeBytes(buf, len);
			if (use_sha256) {
				::SHA256_Update(&content_sha256_buf, buf, len);
			} else {
				::MD5_Update(&content_md5_buf, buf, len);
			}
			length += len;
		}
		// Checksum can only be finalized once, further calls will do nothing so new writes will be reflected in the
		// sum.
		void finalizeChecksum() {
			if (checksumString.empty()) {
				if (use_sha256) {
					std::string sumBytes;
					sumBytes.resize(32);
					::SHA256_Final((unsigned char*)sumBytes.data(), &content_sha256_buf);
					checksumString = base64::encoder::from_string(sumBytes);
					checksumString.resize(checksumString.size() - 1);
				} else {
					std::string sumBytes;
					sumBytes.resize(16);
					::MD5_Final((unsigned char*)sumBytes.data(), &content_md5_buf);
					checksumString = base64::encoder::from_string(sumBytes);
					checksumString.resize(checksumString.size() - 1);
				}
			}
		}

	private:
		bool use_sha256;
		union {
			MD5_CTX content_md5_buf;
			SHA256_CTX content_sha256_buf;
		};
	};

	Future<int> read(void* data, int length, int64_t offset) override { throw file_not_readable(); }

	static Future<Void> write_impl(Reference<AsyncFileBlobStoreWrite> f, const uint8_t* data, int length) {
		Part* p = f->m_parts.back().getPtr();
		// If this write will cause the part to cross the min part size boundary then write to the boundary and start a
		// new part.
		while (p->length + length >= f->m_bstore->knobs.multipart_min_part_size) {
			// Finish off this part
			int finishlen = f->m_bstore->knobs.multipart_min_part_size - p->length;
			p->write((const uint8_t*)data, finishlen);

			// Adjust source buffer args
			length -= finishlen;
			data = (const uint8_t*)data + finishlen;

			// End current part (and start new one)
			co_await f->endCurrentPart(f.getPtr(), true);
			p = f->m_parts.back().getPtr();
		}

		p->write((const uint8_t*)data, length);
	}

	Future<Void> write(void const* data, int length, int64_t offset) override {
		if (offset != m_cursor)
			throw non_sequential_op();
		m_cursor += length;

		return m_error.getFuture() ||
		       write_impl(Reference<AsyncFileBlobStoreWrite>::addRef(this), (const uint8_t*)data, length);
	}

	Future<Void> truncate(int64_t size) override {
		if (size != m_cursor)
			return non_sequential_op();
		return Void();
	}

	static Future<std::string> doPartUpload(AsyncFileBlobStoreWrite* f, Part* p) {
		p->finalizeChecksum();
		std::string upload_id = co_await f->getUploadID();
		std::string etag = co_await f->m_bstore->uploadPart(
		    f->m_bucket, f->m_object, upload_id, p->number, &p->content, p->length, p->checksumString);
		co_return etag;
	}

	static Future<Void> doFinishUpload(AsyncFileBlobStoreWrite* f) {
		// If there is only 1 part then it has not yet been uploaded so just write the whole file at once.
		if (f->m_parts.size() == 1) {
			Reference<Part> part = f->m_parts.back();
			part->finalizeChecksum();
			co_await f->m_bstore->writeEntireFileFromBuffer(
			    f->m_bucket, f->m_object, &part->content, part->length, part->checksumString);
			co_return;
		}

		// There are at least 2 parts.  End the last part (which could be empty)
		co_await f->endCurrentPart(f);

		IBlobStoreEndpoint::MultiPartSetT partSet;
		std::vector<Reference<Part>>::iterator p;

		// Wait for all the parts to be done to get their ETags, populate the partSet required to finish the object
		// upload.
		for (p = f->m_parts.begin(); p != f->m_parts.end(); ++p) {
			std::string tag = co_await (*p)->etag;
			if ((*p)->length > 0) { // The last part might be empty and has to be omitted.
				partSet[(*p)->number] = IBlobStoreEndpoint::PartInfo(tag, (*p)->checksumString);
			}
		}

		// No need to wait for the upload ID here because the above loop waited for all the parts and each part required
		// the upload ID so it is ready
		Optional<std::string> checksumSHA256 = co_await f->m_bstore->finishMultiPartUpload(
		    f->m_bucket, f->m_object, f->m_upload_id.get(), partSet, f->m_cursor);

		// Log the checksum if present - this is just a hash of the multipart structure, not the object content
		if (checksumSHA256.present()) {
			TraceEvent(SevDebug, "AsyncFileBlobStoreMultipartUploadChecksum")
			    .detail("Bucket", f->m_bucket)
			    .detail("Object", f->m_object)
			    .detail("ChecksumSHA256", checksumSHA256.get())
			    .detail("Note", "This is a hash of the multipart structure, not object content");
		}
	}

	// Ready once all data has been sent AND acknowledged from the remote side
	Future<Void> sync() override {
		// Only initiate the finish operation once, and also prevent further writing.
		if (!m_finished.isValid()) {
			m_finished = doFinishUpload(this);
			m_cursor = -1; // Cause future write attempts to fail
		}

		return m_finished;
	}

	//
	// Flush can't really do what the caller would "want" for a blob store file.  The caller would probably notionally
	// want all bytes written to be at least in transit to the blob store, but that is not very feasible.  The blob
	// store has a minimum size requirement for all but the final part, and parts must be sent with a header that
	// specifies their size.  So in the case of a write buffer that does not meet the part minimum size the part could
	// be sent but then if there is any more data written then that part needs to be sent again in its entirety.  So a
	// client that calls flush often could generate far more blob store write traffic than they intend to.
	Future<Void> flush() override { return Void(); }

	Future<int64_t> size() const override { return m_cursor; }

	Future<Void> readZeroCopy(void** data, int* length, int64_t offset) override {
		TraceEvent(SevError, "ReadZeroCopyNotSupported").detail("FileType", "S3BlobStoreWrite");
		return platform_error();
	}
	void releaseZeroCopy(void* data, int length, int64_t offset) override {}

	int64_t debugFD() const override { return -1; }

	~AsyncFileBlobStoreWrite() override {
		m_upload_id.cancel();
		m_finished.cancel();
		m_parts.clear(); // Contains futures
	}

	std::string getFilename() const override { return m_object; }

private:
	Reference<IBlobStoreEndpoint> m_bstore;
	std::string m_bucket;
	std::string m_object;

	int64_t m_cursor;

	Future<std::string> m_upload_id;
	Future<Void> m_finished;
	std::vector<Reference<Part>> m_parts;
	Promise<Void> m_error;
	FlowLock m_concurrentUploads;

	// End the current part and start uploading it, but also wait for a part to finish if too many are in transit.
	static Future<Void> endCurrentPart(AsyncFileBlobStoreWrite* f, bool startNew = false) {
		if (f->m_parts.back()->length == 0)
			co_return;

		// Wait for an upload slot to be available
		co_await f->m_concurrentUploads.take();

		// Do the upload, and if it fails forward errors to m_error and also stop if anything else sends an error to
		// m_error Also, hold a releaser for the concurrent upload slot while all that is going on.
		auto releaser = std::make_shared<FlowLock::Releaser>(f->m_concurrentUploads, 1);
		f->m_parts.back()->etag =
		    holdWhile(releaser, joinErrorGroup(doPartUpload(f, f->m_parts.back().getPtr()), f->m_error));

		// Make a new part to write to
		if (startNew) {
			f->m_parts.push_back(makeReference<Part>(f->m_parts.size() + 1,
			                                         f->m_bstore->knobs.multipart_min_part_size,
			                                         f->m_bstore->knobs.enable_object_integrity_check));
		}
	}

	Future<std::string> getUploadID() {
		if (!m_upload_id.isValid())
			m_upload_id = m_bstore->beginMultiPartUpload(m_bucket, m_object);
		return m_upload_id;
	}

public:
	AsyncFileBlobStoreWrite(Reference<IBlobStoreEndpoint> bstore, std::string bucket, std::string object)
	  : m_bstore(bstore), m_bucket(bucket), m_object(object), m_cursor(0),
	    m_concurrentUploads(bstore->knobs.concurrent_writes_per_file) {

		// Add first part
		m_parts.push_back(makeReference<Part>(
		    1, m_bstore->knobs.multipart_min_part_size, m_bstore->knobs.enable_object_integrity_check));
	}
};

// This class represents a read-only file that lives in a blob store.  It reads using the REST API.
class AsyncFileBlobStoreRead final : public IAsyncFile, public ReferenceCounted<AsyncFileBlobStoreRead> {
public:
	void addref() override { ReferenceCounted<AsyncFileBlobStoreRead>::addref(); }
	void delref() override { ReferenceCounted<AsyncFileBlobStoreRead>::delref(); }

	StringRef getClassName() override { return "AsyncFileBlobStoreRead"_sr; }

	Future<int> read(void* data, int length, int64_t offset) override;

	Future<Void> write(void const* data, int length, int64_t offset) override { throw file_not_writable(); }
	Future<Void> truncate(int64_t size) override { throw file_not_writable(); }

	Future<Void> sync() override { return Void(); }
	Future<Void> flush() override { return Void(); }

	Future<int64_t> size() const override;

	Future<Void> readZeroCopy(void** data, int* length, int64_t offset) override {
		TraceEvent(SevError, "ReadZeroCopyNotSupported").detail("FileType", "S3BlobStoreRead");
		return platform_error();
	}
	void releaseZeroCopy(void* data, int length, int64_t offset) override {}

	int64_t debugFD() const override { return -1; }

	std::string getFilename() const override { return m_object; }

	~AsyncFileBlobStoreRead() override = default;

	Reference<IBlobStoreEndpoint> m_bstore;
	std::string m_bucket;
	std::string m_object;
	mutable Future<int64_t> m_size;

	AsyncFileBlobStoreRead(Reference<IBlobStoreEndpoint> bstore, std::string bucket, std::string object)
	  : m_bstore(bstore), m_bucket(bucket), m_object(object) {}
};