[20240817] XTDB Internals 3 - RemoteBufferPool

/**
 * Implementation for the storage module that persists data remotely within a specified [objectStore],
 * while maintaining a local cache of the working set cache under the [localDiskCache] directory.
 *
 * Any implementer of [ObjectStoreFactory] can be used as the [objectStore]. We currently offer:
 * * AWS S3 (under **xtdb-aws**)
 * * Azure Blob Storage (under **xtdb-azure**)
 * * Google Cloud Storage (under **xtdb-google-cloud**)
 * 
 * Example usage, as part of a node config:
 * ```kotlin
 * Xtdb.openNode {
 *    remoteStorage(
 *       objectStore = objStoreImpl(...) { ... },
 *       localDiskCache = Paths.get("test-path")
 *    ) {
 *       maxCacheEntries = 1024,
 *       maxCacheBytes = 536870912,
 *       maxDiskCachePercentage = 75,
 *       maxDiskCacheBytes = 10737418240
 *    },
 *    ...
 * }
 * ```
 *
 * @property objectStore configuration of the object store to use for remote storage.
 * @property localDiskCache local directory to store the working-set cache in.
 */
@Serializable
@SerialName("!Remote")
data class RemoteStorageFactory(
    val objectStore: ObjectStoreFactory,
    val localDiskCache: Path,
    var maxCacheEntries: Long = 1024,
    var maxCacheBytes: Long = 536870912,
    var maxDiskCachePercentage: Long = 75,
    var maxDiskCacheBytes: Long? = null
) : Factory {

    fun maxCacheEntries(maxCacheEntries: Long) = apply { this.maxCacheEntries = maxCacheEntries }
    fun maxCacheBytes(maxCacheBytes: Long) = apply { this.maxCacheBytes = maxCacheBytes }
    fun maxDiskCachePercentage(maxDiskCachePercentage: Long) = apply { this.maxDiskCachePercentage = maxDiskCachePercentage }
    fun maxDiskCacheBytes(maxDiskCacheBytes: Long) = apply { this.maxDiskCacheBytes = maxDiskCacheBytes }

    override fun openStorage(allocator: BufferAllocator) =
	requiringResolve("xtdb.buffer-pool/open-remote-storage").invoke(allocator, this) as IBufferPool
}

@JvmStatic
fun remoteStorage(objectStore: ObjectStoreFactory, localDiskCachePath: Path) =
    RemoteStorageFactory(objectStore, localDiskCachePath)

@JvmSynthetic
fun Xtdb.Config.remoteStorage(
    objectStore: ObjectStoreFactory,
    localDiskCachePath: Path,
    configure: RemoteStorageFactory.() -> Unit,
) = storage(RemoteStorageFactory(objectStore, localDiskCachePath).also(configure))  

(defmethod xtn/apply-config! ::remote [^Xtdb$Config config _ {:keys [object-store local-disk-cache max-cache-bytes max-cache-entries max-disk-cache-bytes max-disk-cache-percentage]}]
  (.storage config (cond-> (Storage/remoteStorage (let [[tag opts] object-store]  ;; Storage/remoteStorage를 호출하여 기본 원격 저장소 설정을 생성
						    (->object-store-factory tag opts))  ;; 객체 저장소 팩토리로 변환
						  (util/->path local-disk-cache))  ;; local-disk-cache를 path로 변환
		     max-cache-bytes (.maxCacheBytes max-cache-bytes)
		     max-cache-entries (.maxCacheEntries max-cache-entries)
		     max-disk-cache-bytes (.maxDiskCacheBytes max-disk-cache-bytes)
		     max-disk-cache-percentage (.maxDiskCachePercentage max-disk-cache-percentage))))

(defn open-remote-storage ^xtdb.IBufferPool [^BufferAllocator allocator, ^Storage$RemoteStorageFactory factory]
  (util/with-close-on-catch [object-store (.openObjectStore (.getObjectStore factory))]
    (let [^Path local-disk-cache (.getLocalDiskCache factory)
	  local-disk-size-limit (or (.getMaxDiskCacheBytes factory)
				    (calculate-limit-from-percentage-of-disk local-disk-cache (.getMaxDiskCachePercentage factory)))
	  ^AsyncCache local-disk-cache-evictor (->local-disk-cache-evictor local-disk-size-limit local-disk-cache)
	  !evictor-max-weight (atom local-disk-size-limit)]

      ;; Add watcher to max-weight atom - when it changes, update the cache max weight
      (add-watch !evictor-max-weight :update-cache-max-weight
		 (fn [_ _ _ ^long new-size]
		   (let [sync-cache (.synchronous local-disk-cache-evictor)
			 ^Policy$Eviction eviction-policy (.get (.eviction (.policy sync-cache)))]
		     (.setMaximum eviction-policy (max 0 new-size)))))


      (->RemoteBufferPool (.newChildAllocator allocator "buffer-pool" 0 Long/MAX_VALUE)
			  (->memory-buffer-cache (.getMaxCacheBytes factory))
			  local-disk-cache
			  local-disk-cache-evictor
			  object-store
			  !evictor-max-weight))))

(defrecord RemoteBufferPool [allocator
			     ^Cache memory-store
			     ^Path local-disk-cache
			     ^AsyncCache local-disk-cache-evictor
			     ^ObjectStore object-store
			     !evictor-max-weight]
  IBufferPool

  ,,,
  )

(getBuffer [_ k]
    (when k
      (or (cache-get memory-store k)
	  (let [buffer-cache-path (.resolve local-disk-cache k)
		{:keys [file-size ctx]} @(update-evictor-key local-disk-cache-evictor k
							     (fn [^CompletableFuture fut]
							       (-> (if (util/path-exists buffer-cache-path) ;; 버퍼가 로컬 디스크 캐시에 존재하는지 확인
								     ;; 로컬 캐시에 존재하는 경우
								     (-> (or fut (CompletableFuture/completedFuture {:pinned? false}))  ;; 기존 CompletableFuture 사용하던가 새로 생성
									 (.thenApply (fn [{:keys [pinned?]}]
										       (log/tracef "Key %s found in local-disk-cache - returning buffer" k)  ;; 로그를 남기고 버퍼 경로와 이전 고정 상태를 리턴.
										       {:path buffer-cache-path :previously-pinned? pinned?})))
								     ;; 로컬 캐시에 존재하지 않는 경우.
								     ;; 필요한 디렉토리를 생성 
								     (do (util/create-parents buffer-cache-path)
									 (log/debugf "Key %s not found in local-disk-cache, loading from object store" k)
									 ;; 객체 저장소에서 객체를 로드
									 (-> (.getObject object-store k buffer-cache-path)
									     ;; 로드된 path와 previously-pinned? 를 false로 리턴한다.
									     (.thenApply (fn [path]
											   {:path path :previously-pinned? false})))))
								   ;; 이후 공통 로직
								   (.thenApply (fn [{:keys [path previously-pinned?]}]
										 (let [file-size (util/size-on-disk path)  ;; 파일크기 계산
										       nio-buffer (util/->mmap-path path)  ;; 메모리 맵 버퍼 계산
										       buffer-release-fn (fn []  ;; 버퍼 해제 함수 정의, 버퍼가 해제될 때 호출됨.
													   (update-evictor-key local-disk-cache-evictor k
													     (fn [^CompletableFuture fut]
													       (some-> fut
														       (.thenApply (fn [{:keys [pinned? file-size]}]
																     ;; buffer 를 release 할때는 !evictor-max-weight의 크기를 키운다.
																     ;; pinned인 경우만 빼기를 했을테니까 원래대로 되돌리는 것.(덧셈으로)
																     (when pinned?
																       (swap! !evictor-max-weight + file-size))
																     ;; 고정 상태를 해제하고 필요한 경우 전체 무게를 조정함.
																     {:pinned? false :file-size file-size}))))))
										       ;; ArrowBuf 생성
										       create-arrow-buf #(util/->arrow-buf-view allocator nio-buffer buffer-release-fn)
										       ;; 메모리 캐시에 저장.
										       buf (cache-compute memory-store k create-arrow-buf)]
										   ;; 해당 버퍼를 pinned? true 로 변경, 파일 크기, 버퍼, previously-pinned? 정보 포함.
										   {:pinned? true :file-size file-size :ctx {:buf buf :previously-pinned? previously-pinned?}}))))))
		{:keys [previously-pinned? buf]} ctx]

	    (when-not previously-pinned?
	      (swap! !evictor-max-weight - file-size))
	    buf))))

(putObject [_ k buffer]
  @(.putObject object-store k buffer))

package xtdb.api.storage

import java.nio.ByteBuffer
import java.nio.file.Path
import java.util.concurrent.CompletableFuture

interface ObjectStore : AutoCloseable {
    /**
     * Asynchronously returns the given object in a ByteBuffer.
     *
     * If the object doesn't exist, the CompletableFuture completes with an IllegalStateException.
     */
    fun getObject(k: Path): CompletableFuture<ByteBuffer>

    /**
     * Asynchronously returns a specified range of bytes from the object in a ByteBuffer.
     *
     * If the object doesn't exist, the CompletableFuture completes with an IllegalStateException.
     */
    fun getObjectRange(k: Path, start: Long, len: Long): CompletableFuture<ByteBuffer>

    /**
     * Asynchronously writes the object to the given path.
     *
     * If the object doesn't exist, the CompletableFuture completes with an IllegalStateException.
     */
    fun getObject(k: Path, outPath: Path): CompletableFuture<Path>

    /**
     * Stores an object in the object store.
     */
    fun putObject(k: Path, buf: ByteBuffer): CompletableFuture<*>

    /**
     * Recursively lists all objects in the object store.
     *
     * Objects are returned in lexicographic order of their path names.
     */
    fun listAllObjects(): Iterable<Path>

    /**
     * Lists objects directly within the specified directory in the object store.
     *
     * Objects are returned in lexicographic order of their path names.
     */
    fun listObjects(dir: Path): Iterable<Path>

    /**
     * Deletes the object with the given path from the object store.
     */
    fun deleteObject(k: Path): CompletableFuture<*>

    override fun close() {
    }
}

interface ObjectStoreFactory {
    fun openObjectStore(): ObjectStore
}  

(openArrowWriter [_ k rel]
    (let [tmp-path (create-tmp-path local-disk-cache)]
      (util/with-close-on-catch [file-ch (util/->file-channel tmp-path util/write-truncate-open-opts)
				 unl (.startUnload rel file-ch)]

	(reify ArrowWriter
	  (writeBatch [_] (.writeBatch unl))

	  (end [_]
	    (.endFile unl)
	    (.close file-ch)

	    (upload-arrow-file allocator object-store k tmp-path)

	    (let [buffer-cache-path (.resolve local-disk-cache k)]
	      (update-evictor-key local-disk-cache-evictor k
				  (fn [^CompletableFuture fut]
				    (-> (or fut (CompletableFuture/completedFuture {:pinned? false}))
					(.thenApply (fn [{:keys [pinned?]}]
						      (log/tracef "Writing arrow file, %s, to local disk cache" k)
						      (util/create-parents buffer-cache-path)
						      ;; see #2847
						      (util/atomic-move tmp-path buffer-cache-path)
						      {:pinned? pinned? :file-size (util/size-on-disk buffer-cache-path)})))))))

	  (close [_]
	    (util/close unl)
	    (when (.isOpen file-ch)
	      (.close file-ch)))))))