open Equinox.Core

open FSharp.Control

open FsCodec

open Microsoft.Azure.Cosmos

open Serilog

open System

open System.Collections.Generic

open System.Text.Json

type EventBody = JsonElement

type Event = // TODO for STJ v5: All fields required unless explicitly optional

{ /// Creation datetime (as opposed to system-defined _lastUpdated which is touched by triggers, replication etc.)

t: DateTimeOffset // ISO 8601

/// The Case (Event Type); used to drive deserialization

c: string // required

/// Event body, as UTF-8 encoded json ready to be injected into the Json being rendered for CosmosDB

d: EventBody // TODO for STJ v5: Required, but can be null so Nullary cases can work

/// Optional metadata, as UTF-8 encoded json, ready to emit directly

m: EventBody // TODO for STJ v5: Optional, not serialized if missing

/// Optional correlationId

correlationId: string // TODO for STJ v5: Optional, not serialized if missing

/// Optional causationId

causationId: string } // TODO for STJ v5: Optional, not serialized if missing

interface IEventData<EventBody> with

member x.EventType = x.c

member x.Data = x.d

member x.Meta = x.m

member _.EventId = Guid.Empty

member x.CorrelationId = x.correlationId

member x.CausationId = x.causationId

member x.Timestamp = x.t

type Batch = // TODO for STJ v5: All fields required unless explicitly optional

{ /// CosmosDB-mandated Partition Key, must be maintained within the document

/// Not actually required if running in single partition mode, but for simplicity, we always write it

p: string // "{streamName}" TODO for STJ v5: Optional, not requested in queries

/// CosmosDB-mandated unique row key; needs to be unique within any partition it is maintained; must be string

/// At the present time, one can't perform an ORDER BY on this field, hence we also have i shadowing it

/// NB Tip uses a well known value here while it's actively 'open'

id: string // "{index}"

/// When we read, we need to capture the value so we can retain it for caching purposes

/// NB this is not relevant to fill in when we pass it to the writing stored procedure

/// as it will do: 1. read 2. merge 3. write merged version contingent on the _etag not having changed

_etag: string // TODO for STJ v5: Optional, not serialized if missing

/// base 'i' value for the Events held herein

i: int64 // {index}

// `i` value for successor batch (to facilitate identifying which Batch a given startPos is within)

n: int64 // {index}

/// The Domain Events (as opposed to Unfolded Events, see Tip) at this offset in the stream

e: Event[] }

static member internal PartitionKeyField = "p"

// As one cannot sort by the (mandatory) `id` field, we have an indexed `i` field for sort and range query use

// NB its critical to also index the nominated PartitionKey field as RU costs increase (things degrade to scans) if you don't

// TOCONSIDER: indexing strategy was developed and tuned before composite key extensions in ~2021, which might potentially be more efficient

// a decent attempt at https://github.com/jet/equinox/issues/274 failed, but not 100% sure it's fundamentally impossible/wrong

static member internal IndexedPaths = [| Batch.PartitionKeyField; "i"; "n" |] |> Array.map (fun k -> $"/%s{k}/?")

type Unfold =

{ /// Base: Stream Position (Version) of State from which this Unfold Event was generated. An unfold from State Version 1 is i=1 and includes event i=1

i: int64

/// Generation datetime

t: DateTimeOffset // ISO 8601 // Not written by versions <= 2.0.0-rc9

/// The Case (Event Type) of this compaction/snapshot, used to drive deserialization

c: string // required

/// Event body - Json -> Deflate -> Base64 -> JsonElement

[<Serialization.JsonConverter(typeof<JsonCompressedBase64Converter>)>]

d: EventBody // required

/// Optional metadata, same encoding as `d` (can be null; not written if missing)

[<Serialization.JsonConverter(typeof<JsonCompressedBase64Converter>)>]

m: EventBody } // TODO for STJ v5: Optional, not serialized if missing

// Arrays are not indexed by default. 1. enable filtering by `c`ase 2. index uncompressed fields within unfolds for filtering

static member internal IndexedPaths = [| "/u/[]/c/?"; "/u/[]/d/*" |]

type Tip = // TODO for STJ v5: All fields required unless explicitly optional

{ /// Partition key, as per Batch

p: string // "{streamName}" TODO for STJ v5: Optional, not requested in queries

/// Document Id within partition, as per Batch

id: string // "{-1}" - Well known IdConstant used while this remains the pending batch

/// When we read, we need to capture the value so we can retain it for caching purposes

/// NB this is not relevant to fill in when we pass it to the writing stored procedure

/// as it will do: 1. read 2. merge 3. write merged version contingent on the _etag not having changed

_etag: string // TODO for STJ v5: Optional, not serialized if missing

/// base 'i' value for the Events held herein

i: int64

/// `i` value for successor batch (to facilitate identifying which Batch a given startPos is within)

n: int64 // {index}

/// Domain Events, will eventually move out to a Batch

e: Event[]

/// Compaction/Snapshot/Projection events - owned and managed by the sync stored proc

u: Unfold[] }

static member internal WellKnownDocumentId = "-1"

type Position = { index: int64; etag: string option }

let fromEtagOnly (value: string) = { index = -2; etag = Some value }

let fromEtagAndIndex (etag, n) = { index = n; etag = Some etag }

/// NB a token without an etag is inefficient compared to fromEtagAndIndex, so paths to this should be minimized

let fromIndex (x: int64) = { index = x; etag = None }

/// If we have strong reason to suspect a stream is empty, we won't have an etag (and Writer Stored Procedure special cases this)

let fromKnownEmpty = fromIndex 0L

/// Blind append mode

let fromAppendAtEnd = fromIndex -1L // sic - needs to yield value -1 to trigger stored proc logic

/// Sentinel value we assign so we can reject attempts to sync without having known the context

let readOnly = { index = -2L; etag = None }

let isReadOnly (x: Position) = x.index = -2L && Option.isNone x.etag

/// If we encounter the tip (id=-1) item/document, we're interested in its etag so we can re-sync for 1 RU

let tryFromBatch (x: Batch) =

if x.id <> Tip.WellKnownDocumentId then None

else Some { index = x.n; etag = Option.ofObj x._etag }

type Direction = Forward | Backward override this.ToString() = match this with Forward -> "Forward" | Backward -> "Backward"

type internal Enum private () =

static member Events(i: int64, e: Event[], ?minIndex, ?maxIndex): ITimelineEvent<EventBody> seq = seq {

let indexMin, indexMax = defaultArg minIndex 0L, defaultArg maxIndex Int64.MaxValue

for offset in 0..e.Length-1 do

let index = i + int64 offset

// If we're loading from a nominated position, we need to discard items in the batch before/after the start on the start page

if index >= indexMin && index < indexMax then

let x = e[offset]

yield FsCodec.Core.TimelineEvent.Create(index, x.c, x.d, x.m, Guid.Empty, x.correlationId, x.causationId, x.t) }

static member internal Events(t: Tip, ?minIndex, ?maxIndex): ITimelineEvent<EventBody> seq =

Enum.Events(t.i, t.e, ?minIndex = minIndex, ?maxIndex = maxIndex)

static member internal Events(b: Batch, ?minIndex, ?maxIndex) =

Enum.Events(b.i, b.e, ?minIndex = minIndex, ?maxIndex = maxIndex)

static member Unfolds(xs: Unfold[]): ITimelineEvent<EventBody> seq = seq {

for x in xs -> FsCodec.Core.TimelineEvent.Create(x.i, x.c, x.d, x.m, Guid.Empty, null, null, x.t, isUnfold = true) }

static member EventsAndUnfolds(x: Tip, ?minIndex, ?maxIndex): ITimelineEvent<EventBody> seq =

Enum.Events(x, ?minIndex = minIndex, ?maxIndex = maxIndex)

|> Seq.append (Enum.Unfolds x.u)

// where Index is equal, unfolds get delivered after the events so the fold semantics can be 'idempotent'

|> Seq.sortBy (fun x -> x.Index, x.IsUnfold)

type IRetryPolicy = abstract member Execute: (int -> CancellationToken -> Task<'T>) -> Task<'T>

/// <summary>Name of Property used for <c>Metric</c> in <c>LogEvent</c>s.</summary>

let [<Literal>] PropertyTag = "cosmosEvt"

[<NoEquality; NoComparison>]

type Measurement =

{ database: string; container: string; stream: string

interval: StopwatchInterval; bytes: int; count: int; ru: float }

member x.Category = x.stream |> StreamName.Internal.trust |> StreamName.Category.ofStreamName

[<RequireQualifiedAccess; NoEquality; NoComparison>]

type Metric =

/// Individual read request for the Tip

| Tip of Measurement

/// Individual read request for the Tip, not found

| TipNotFound of Measurement

/// Tip read with Single RU Request Charge due to correct use of etag in cache

| TipNotModified of Measurement

/// Summarizes a set of Responses for a given Read request

| Query of Direction * responses: int * Measurement

/// Individual read request in a Batch

/// Charges are rolled up into Query Metric (so do not double count)

| QueryResponse of Direction * Measurement

| SyncSuccess of Measurement

| SyncResync of Measurement

| SyncConflict of Measurement

/// Summarizes outcome of request to trim batches from head of a stream and events in Tip

/// Count in Measurement is number of batches (documents) deleted

/// Bytes in Measurement is number of events deleted

| Prune of responsesHandled: int * Measurement

/// Handled response from listing of batches in a stream

/// Charges are rolled up into the Prune Metric (so do not double count)

| PruneResponse of Measurement

/// Deleted an individual Batch

| Delete of Measurement

/// Trimmed the Tip

| Trim of Measurement

let [<return: Struct>] (|MetricEvent|_|) (logEvent: Serilog.Events.LogEvent): Metric voption =

let mutable p = Unchecked.defaultof<_>

logEvent.Properties.TryGetValue(PropertyTag, &p) |> ignore

match p with Log.ScalarValue (:? Metric as e) -> ValueSome e | _ -> ValueNone

/// Attach a property to the captured event record to hold the metric information

let internal event (value: Metric) = Internal.Log.withScalarProperty PropertyTag value

let internal prop name value (log: ILogger) = log.ForContext(name, value)

let internal propData name (events: #IEventData<EventBody> seq) (log: ILogger) =

let render (body: EventBody) = body.GetRawText()

let items = seq { for e in events do yield sprintf "{\"%s\": %s}" e.EventType (render e.Data) }

log.ForContext(name, sprintf "[%s]" (String.concat ",\n\r" items))

let internal propEvents = propData "events"

let internal propDataUnfolds = Enum.Unfolds >> propData "unfolds"

let internal propStartPos (value: Position) log = prop "startPos" value.index log

let internal propStartEtag (value: Position) log = prop "startEtag" value.etag log

let withLoggedRetries<'t> (retryPolicy: IRetryPolicy option) (contextLabel: string) (f: ILogger -> CancellationToken -> Task<'t>) log ct: Task<'t> =

match retryPolicy with

| None -> f log ct

| Some retryPolicy ->

let withLoggingContextWrapping count =

let log = if count = 1 then log else log |> prop contextLabel count

f log

retryPolicy.Execute withLoggingContextWrapping

let internal (|BlobLen|) (x: EventBody) = if x.ValueKind = JsonValueKind.Null then 0 else x.GetRawText().Length

let internal eventLen (x: #IEventData<_>) = let BlobLen bytes, BlobLen metaBytes = x.Data, x.Meta in bytes + metaBytes + 80

let internal batchLen = Seq.sumBy eventLen

[<RequireQualifiedAccess>]

type Operation = Tip | Tip404 | Tip304 | Query | Write | Resync | Conflict | Prune | Delete | Trim

let (|Op|QueryRes|PruneRes|) = function

| Metric.Tip s -> Op (Operation.Tip, s)

| Metric.TipNotFound s -> Op (Operation.Tip404, s)

| Metric.TipNotModified s -> Op (Operation.Tip304, s)

| Metric.Query (_, _, s) -> Op (Operation.Query, s)

| Metric.QueryResponse (direction, s) -> QueryRes (direction, s)

| Metric.SyncSuccess s -> Op (Operation.Write, s)

| Metric.SyncResync s -> Op (Operation.Resync, s)

| Metric.SyncConflict s -> Op (Operation.Conflict, s)

| Metric.Prune (_, s) -> Op (Operation.Prune, s)

| Metric.PruneResponse s -> PruneRes s

| Metric.Delete s -> Op (Operation.Delete, s)

| Metric.Trim s -> Op (Operation.Trim, s)

/// NB Caveat emptor; this is subject to unlimited change without the major version changing - while the `dotnet-templates` repo will be kept in step, and

/// the ChangeLog will mention changes, it's critical to not assume that the presence or nature of these helpers be considered stable

module InternalMetrics =

module Stats =

type internal Counter =

{ mutable minRu: float; mutable maxRu: float; mutable rux100: int64; mutable count: int64; mutable ms: int64 }

static member Create() = { minRu = Double.MaxValue; maxRu = 0; rux100 = 0L; count = 0L; ms = 0L }

member x.Ingest(ms, ru) =

Interlocked.Increment(&x.count) |> ignore

Interlocked.Add(&x.rux100, int64 (ru * 100.)) |> ignore

if ru < x.minRu then Interlocked.CompareExchange(&x.minRu, ru, x.minRu) |> ignore

if ru > x.maxRu then Interlocked.CompareExchange(&x.maxRu, ru, x.maxRu) |> ignore

Interlocked.Add(&x.ms, ms) |> ignore

type internal Counters() =

let buckets = System.Collections.Concurrent.ConcurrentDictionary<string, Counter>()

let create (_name: string) = Counter.Create()

member _.Ingest(bucket, ms, ru) = buckets.GetOrAdd(bucket, create).Ingest(ms, ru)

member _.Buckets = buckets.Keys

member _.TryBucket bucket = match buckets.TryGetValue bucket with true, t -> Some t | false, _ -> None

type Epoch() =

let epoch = System.Diagnostics.Stopwatch.StartNew()

member val internal Tip = Counters() with get, set

member val internal Read = Counters() with get, set

member val internal Write = Counters() with get, set

member val internal Resync = Counters() with get, set

member val internal Conflict = Counters() with get, set

member val internal Prune = Counters() with get, set

member val internal Delete = Counters() with get, set

member val internal Trim = Counters() with get, set

member _.Stop() = epoch.Stop()

member _.Elapsed = epoch.Elapsed

type LogSink(categorize) =

let bucket (x: Measurement) = if categorize then $"{x.container}/{x.Category}" else x.container

let (|BucketMsRu|) ({ interval = i; ru = ru } as m) = bucket m, int64 i.ElapsedMilliseconds, ru

static let mutable epoch = Epoch()

new() = LogSink(false)

static member Restart() =

let fresh = Epoch()

let outgoing = Interlocked.Exchange(&epoch, fresh)

outgoing.Stop()

outgoing

interface Serilog.Core.ILogEventSink with

member _.Emit logEvent =

match logEvent with

| MetricEvent cm ->

match cm with

| Op ((Operation.Tip | Operation.Tip404 | Operation.Tip304), BucketMsRu m) ->

epoch.Tip.Ingest m

| Op (Operation.Query, BucketMsRu m) -> epoch.Read.Ingest m

| QueryRes (_direction, _) -> ()

| Op (Operation.Write, BucketMsRu m) -> epoch.Write.Ingest m

| Op (Operation.Conflict, BucketMsRu m) -> epoch.Conflict.Ingest m

| Op (Operation.Resync, BucketMsRu m) -> epoch.Resync.Ingest m

| Op (Operation.Prune, BucketMsRu m) -> epoch.Prune.Ingest m

| PruneRes _ -> ()

| Op (Operation.Delete, BucketMsRu m) -> epoch.Delete.Ingest m

| Op (Operation.Trim, BucketMsRu m) -> epoch.Trim.Ingest m

| _ -> ()

/// Relies on feeding of metrics from Log through to Stats.LogSink

/// Use Stats.LogSink.Restart() to reset the start point (and stats) where relevant

let dump (log: ILogger) =

let res = Stats.LogSink.Restart()

let stats =

[|nameof res.Tip, res.Tip

nameof res.Read, res.Read

nameof res.Write, res.Write

nameof res.Resync, res.Resync

nameof res.Conflict, res.Conflict

nameof res.Prune, res.Prune

nameof res.Delete, res.Delete

nameof res.Trim, res.Trim |]

let isRead = function nameof res.Tip | nameof res.Read | nameof res.Prune -> true | _ -> false

let buckets = stats |> Seq.collect (fun (_n, stat) -> stat.Buckets) |> Seq.distinct |> Seq.sort |> Seq.toArray

if Array.isEmpty buckets then () else

let maxBucketLen = buckets |> Seq.map _.Length |> Seq.max

let duration = res.Elapsed.TotalSeconds

let mutable prevCat, catR, catW, catRRu, catWRu = null, 0L, 0L, 0., 0.

let inline rps count = if duration = 0 then 0L else float count/duration |> int64

let inline ups ru = if duration = 0 then 0. else ru/duration

let logOnCatChange (cat: string) =

if prevCat = null then prevCat <- cat

elif prevCat = cat then ()

else

let reqs = catR + catW

log.Information("{bucket} {count,6}r @ {rps,5:f0} r/s {rups,5:f0} RU/s ({rrups:f0}{r:l}/{wrups:f0}{w:l})",

prevCat.PadRight maxBucketLen, reqs, rps reqs, ups (catRRu + catWRu), ups catRRu, "R", ups catWRu, "W")

catR <- 0; catRRu <- 0; catW <- 0; catWRu <- 0; prevCat <- cat

for bucket in buckets do

let group = match bucket.IndexOf '/' with -1 -> bucket | i -> bucket.Substring(0, i)

group |> logOnCatChange

for act, counts in stats do

match counts.TryBucket bucket with

| Some stat when stat.count <> 0L ->

let ru = float stat.rux100 / 100.

let rut = if isRead act then catR <- catR + stat.count; catRRu <- catRRu + ru; "R"

else catW <- catW + stat.count; catWRu <- catWRu + ru; "W"

let inline avg x = x / float stat.count

log.Information("{bucket} {act,-8}{count,6}r {minRu,5:f1}-{maxRu,3:f0} {rut:l}RU{lat,5:f0} ms @ {rps,5:f0} r/s {rups,5:f0} RU/s Σ {ru,7:f0} avg={avgRu,4:f1}",

bucket.PadRight maxBucketLen, act, stat.count, stat.minRu, stat.maxRu, rut, avg (float stat.ms), rps stat.count, ups ru, ru, avg ru)

| _ -> ()

null |> logOnCatChange

type ReadResult<'T> = Found of 'T | NotFound | NotModified

type Container with

member private container.DeserializeResponseBody<'T>(rm: ResponseMessage): 'T =

rm.EnsureSuccessStatusCode().Content

|> container.Database.Client.ClientOptions.Serializer.FromStream<'T>

member container.TryReadItem(partitionKey: PartitionKey, id: string, ct, ?options: ItemRequestOptions): Task<float * ReadResult<'T>> = task {

use! rm = container.ReadItemStreamAsync(id, partitionKey, requestOptions = Option.toObj options, cancellationToken = ct)

return rm.Headers.RequestCharge, rm.StatusCode |> function

| System.Net.HttpStatusCode.NotFound -> NotFound

| System.Net.HttpStatusCode.NotModified -> NotModified

| _ -> container.DeserializeResponseBody<'T>(rm) |> Found }

type SyncResponse = { etag: string; n: int64; conflicts: Unfold[]; e: Event[] }

let [<Literal>] name = "EquinoxEventsInTip4" // NB need to rename/number for any breaking change

let [<Literal>] body = """

type internal SyncExp private = Version of int64 | Etag of string | Any

let fromVersion i =

if i < 0L then raise <| ArgumentOutOfRangeException(nameof i, i, "must be >= 0")

SyncExp.Version i

let fromVersionOrAppendAtEnd = function -1L -> SyncExp.Any | i -> fromVersion i

let fromEtag etag = SyncExp.Etag etag

// NB don't nest in a private module, or serialization will fail miserably ;)

[<RequireQualifiedAccess; NoEquality; NoComparison>]

type Result =

| Written of Position

| Conflict of Position * events: ITimelineEvent<EventBody>[]

| ConflictUnknown of Position

let private run (container: Container, stream: string) (maxEventsInTip, maxStringifyLen) (exp, req: Tip, ct)

: Task<float*Result> = task {

let ep =

match exp with

| SyncExp.Version ev -> Position.fromIndex ev

| SyncExp.Etag et -> Position.fromEtagOnly et

| SyncExp.Any -> Position.fromAppendAtEnd

let args = [| box req; box ep.index; box (Option.toObj ep.etag); box maxEventsInTip; box maxStringifyLen |]

let! (res: Scripts.StoredProcedureExecuteResponse<SyncResponse>) =

container.Scripts.ExecuteStoredProcedureAsync<SyncResponse>(SyncStoredProc.name, PartitionKey stream, args, cancellationToken = ct)

let newPos = { index = res.Resource.n; etag = Option.ofObj res.Resource.etag }

match res.Resource.conflicts with

| null -> return res.RequestCharge, Result.Written newPos

// ConflictUnknown is to be yielded if we believe querying is going to be necessary (as there are no unfolds, and no relevant events in the Tip)

| [||] when res.Resource.e.Length = 0 && newPos.index > ep.index ->

return res.RequestCharge, Result.ConflictUnknown newPos

| unfolds -> // stored proc only returns events and unfolds with index >= req.i - no need to trim to a minIndex

let events = (Enum.Events(ep.index, res.Resource.e), Enum.Unfolds unfolds) ||> Seq.append |> Array.ofSeq

return res.RequestCharge, Result.Conflict (newPos, events) }

let private logged (container, stream) (maxEventsInTip, maxStringifyLen) (exp: SyncExp, req: Tip) (log: ILogger) ct: Task<Result> = task {

let! t, (ru, result) = (fun ct -> run (container, stream) (maxEventsInTip, maxStringifyLen) (exp, req, ct)) |> Stopwatch.time ct

let verbose = log.IsEnabled Serilog.Events.LogEventLevel.Debug

let count, bytes = req.e.Length, if verbose then Enum.Events req |> Log.batchLen else 0

let log =

let inline mkMetric ru: Log.Measurement =

{ database = container.Database.Id; container = container.Id; stream = stream; interval = t; bytes = bytes; count = count; ru = ru }

let inline propConflict log = log |> Log.prop "conflict" true |> Log.prop "eventTypes" (Seq.truncate 5 (seq { for x in req.e -> x.c }))

if verbose then log |> Log.propEvents (Enum.Events req) |> Log.propDataUnfolds req.u else log

|> match exp with

| SyncExp.Etag et -> Log.prop "expectedEtag" et

| SyncExp.Version ev -> Log.prop "expectedVersion" ev

| SyncExp.Any -> Log.prop "expectedVersion" -1

|> match result with

| Result.Written pos ->

Log.prop "nextExpectedVersion" pos >> Log.event (Log.Metric.SyncSuccess (mkMetric ru))

| Result.ConflictUnknown pos' ->

Log.prop "nextExpectedVersion" pos' >> propConflict >> Log.event (Log.Metric.SyncConflict (mkMetric ru))

| Result.Conflict (pos', xs) ->

if verbose then Log.propData "conflicts" xs else id

>> Log.prop "nextExpectedVersion" pos' >> propConflict >> Log.event (Log.Metric.SyncResync (mkMetric ru))

log.Information("EqxCosmos {action:l} {stream} {count}+{ucount} {ms:f1}ms {ru}RU {bytes:n0}b {exp}",

"Sync", stream, count, req.u.Length, t.ElapsedMilliseconds, ru, bytes, exp)

return result }

let batch (log: ILogger) (retryPolicy, maxEventsInTip, maxStringifyLen) containerStream expBatch ct: Task<Result> =

let call = logged containerStream (maxEventsInTip, maxStringifyLen) expBatch

Log.withLoggedRetries retryPolicy "writeAttempt" call log ct

let private mkEvent (e: IEventData<_>) =

{ t = e.Timestamp

c = e.EventType; d = JsonElement.undefinedToNull e.Data; m = JsonElement.undefinedToNull e.Meta

correlationId = e.CorrelationId; causationId = e.CausationId }

let mkBatch (stream: string) (events: IEventData<_>[]) unfolds: Tip =

{ p = stream; id = Tip.WellKnownDocumentId; n = -1L(*Server-managed*); i = -1L(*Server-managed*); _etag = null

e = Array.map mkEvent events; u = unfolds }

let mkUnfold baseIndex (compressor, x: IEventData<'F>): Unfold =

{ i = baseIndex; t = x.Timestamp

c = x.EventType; d = compressor x.Data; m = compressor x.Meta }

type [<RequireQualifiedAccess>] Throughput = Manual of rus: int | Autoscale of maxRus: int

type [<RequireQualifiedAccess>] Provisioning = Container of Throughput | Database of Throughput | Serverless

let private (|ThroughputProperties|) = function

| Throughput.Manual rus -> ThroughputProperties.CreateManualThroughput(rus)

| Throughput.Autoscale maxRus -> ThroughputProperties.CreateAutoscaleThroughput(maxRus)

let private createDatabaseIfNotExists (client: CosmosClient) maybeTp dName = async {

let! r = Async.call (fun ct -> client.CreateDatabaseIfNotExistsAsync(dName, throughputProperties = Option.toObj maybeTp, cancellationToken = ct))

return r.Database }

let private createOrProvisionDatabase (client: CosmosClient) dName = function

| Provisioning.Container _ | Provisioning.Serverless -> createDatabaseIfNotExists client None dName

| Provisioning.Database (ThroughputProperties tp) -> async {

let! d = createDatabaseIfNotExists client (Some tp) dName

do! Async.call (fun ct -> d.ReplaceThroughputAsync(tp, cancellationToken = ct)) |> Async.Ignore

return d }

let private createContainerIfNotExists (d: Database) cp maybeTp = async {

let! r = Async.call (fun ct -> d.CreateContainerIfNotExistsAsync(cp, throughputProperties = Option.toObj maybeTp, cancellationToken = ct))

let existed = r.StatusCode = Net.HttpStatusCode.OK

if existed then

do! Async.call (fun ct -> r.Container.ReplaceContainerAsync(cp, cancellationToken = ct)) |> Async.Ignore

return r.Container, existed }

let private createOrProvisionContainer (d: Database) (cName, pkPath, customizeContainer) mode =

let cp = ContainerProperties(id = cName, partitionKeyPath = pkPath)

customizeContainer cp

match mode with

| Provisioning.Database _ | Provisioning.Serverless -> async {

let! c, _existed = createContainerIfNotExists d cp None

return c }

| Provisioning.Container (ThroughputProperties throughput) -> async {

let! c, existed = createContainerIfNotExists d cp (Some throughput)

if existed then do! Async.call (fun ct -> c.ReplaceThroughputAsync(throughput, cancellationToken = ct)) |> Async.Ignore

return c }

let private createStoredProcIfNotExists (c: Container) (name, body) ct: Task<float> = task {

try let! r = c.Scripts.CreateStoredProcedureAsync(Scripts.StoredProcedureProperties(id = name, body = body), cancellationToken = ct)

return r.RequestCharge

with :? CosmosException as ce when ce.StatusCode = System.Net.HttpStatusCode.Conflict -> return ce.RequestCharge }

let private applyBatchAndTipContainerProperties indexUnfolds (cp: ContainerProperties) =

cp.IndexingPolicy.IndexingMode <- IndexingMode.Consistent

cp.IndexingPolicy.Automatic <- true

// We specify fields on whitelist basis; generic querying is inapplicable, and indexing is far from free (write RU and latency)

// Given how long and variable the blacklist would be, we whitelist instead

cp.IndexingPolicy.ExcludedPaths.Add(ExcludedPath(Path="/*"))

for p in [| yield! Batch.IndexedPaths; if indexUnfolds then yield! Unfold.IndexedPaths |] do

cp.IndexingPolicy.IncludedPaths.Add(IncludedPath(Path = p))

let createSyncStoredProcIfNotExists (log: ILogger option) container ct = task {

let! t, ru = createStoredProcIfNotExists container (SyncStoredProc.name, SyncStoredProc.body) |> Stopwatch.time ct

match log with

| None -> ()

| Some log -> log.Information("Created stored procedure {procName} in {ms:f1}ms {ru}RU", SyncStoredProc.name, t.ElapsedMilliseconds, ru) }

let init log (client: CosmosClient) (dName, cName) mode indexUnfolds skipStoredProc ct = task {

let! d = createOrProvisionDatabase client dName mode

let! c = createOrProvisionContainer d (cName, $"/%s{Batch.PartitionKeyField}", applyBatchAndTipContainerProperties indexUnfolds) mode

if not skipStoredProc then

do! createSyncStoredProcIfNotExists (Some log) c ct }

let private applyAuxContainerProperties (cp: ContainerProperties) =

// TL;DR no indexing of any kind; see https://github.com/Azure/azure-documentdb-changefeedprocessor-dotnet/issues/142

cp.IndexingPolicy.Automatic <- false

cp.IndexingPolicy.IndexingMode <- IndexingMode.None

let initAux (client: CosmosClient) (dName, cName) mode = async {

let! d = createOrProvisionDatabase client dName mode

return! createOrProvisionContainer d (cName, "/id", applyAuxContainerProperties) mode } // per Cosmos team, Partition Key must be "/id"

/// Per Container, we need to ensure the stored procedure has been created exactly once (per process lifetime)

type internal ContainerInitializerGuard(container: Container, ?initContainer: Container -> CancellationToken -> Task<unit>) =

let initGuard = initContainer |> Option.map (fun init -> TaskCell<unit>(init container))

member val Container = container

/// Coordinates max of one in flight call to the init logic, retrying on next request if it fails. Calls after it has succeeded noop

member _.Initialize(ct): System.Threading.Tasks.ValueTask =

match initGuard with

| Some g when not (g.IsValid()) -> g.Await(ct) |> ValueTask.ofTask |> ValueTask.ignore

| _ -> System.Threading.Tasks.ValueTask.CompletedTask

let private get (container: Container, stream: string) (maybePos: Position option) ct =

let ro = match maybePos with Some { etag = Some etag } -> ItemRequestOptions(IfNoneMatchEtag = etag) |> Some | _ -> None

container.TryReadItem<Tip>(PartitionKey stream, Tip.WellKnownDocumentId, ct, ?options = ro)

let private loggedGet (get: Container * string -> Position option -> CancellationToken -> Task<_>) (container, stream) (maybePos: Position option) (log: ILogger) ct = task {

let log = log |> Log.prop "stream" stream

let! t, (ru, res: ReadResult<Tip>) = get (container, stream) maybePos |> Stopwatch.time ct

let verbose = log.IsEnabled Events.LogEventLevel.Debug

let log bytes count (f: Log.Measurement -> _) = log |> Log.event (f { database = container.Database.Id; container = container.Id; stream = stream; interval = t; bytes = bytes; count = count; ru = ru })

match res with

| ReadResult.NotModified ->

(log 0 0 Log.Metric.TipNotModified).Information("EqxCosmos {action:l} {stream} {res} {ms:f1}ms {ru}RU", "Tip", stream, 304, t.ElapsedMilliseconds, ru)

| ReadResult.NotFound ->

(log 0 0 Log.Metric.TipNotFound).Information("EqxCosmos {action:l} {stream} {res} {ms:f1}ms {ru}RU", "Tip", stream, 404, t.ElapsedMilliseconds, ru)

| ReadResult.Found tip ->

let log =

let count, bytes = tip.u.Length, if verbose then Enum.Unfolds tip.u |> Log.batchLen else 0

log bytes count Log.Metric.Tip

let log = if verbose then log |> Log.propDataUnfolds tip.u else log

let log = match maybePos with Some p -> log |> Log.propStartPos p |> Log.propStartEtag p | None -> log

let log = log |> Log.prop "_etag" tip._etag |> Log.prop "n" tip.n

log.Information("EqxCosmos {action:l} {stream} {res} {ms:f1}ms {ru}RU", "Tip", stream, 200, t.ElapsedMilliseconds, ru)

return ru, res }

type [<RequireQualifiedAccess; NoComparison; NoEquality>] Result = NotModified | NotFound | Found of Position * i: int64 * ITimelineEvent<EventBody>[]

/// `pos` being Some implies that the caller holds a cached value and hence is ready to deal with Result.NotModified

let tryLoad (log: ILogger) retryPolicy containerStream (maybePos: Position option, maxIndex) ct: Task<Result> = task {

let! _rc, res = Log.withLoggedRetries retryPolicy "readAttempt" (loggedGet get containerStream maybePos) log ct

match res with

| ReadResult.NotModified -> return Result.NotModified

| ReadResult.NotFound -> return Result.NotFound

| ReadResult.Found tip ->

let minIndex = maybePos |> Option.map _.index

return Result.Found (Position.fromEtagAndIndex (tip._etag, tip.n), tip.i, Enum.EventsAndUnfolds(tip, ?maxIndex = maxIndex, ?minIndex = minIndex) |> Array.ofSeq) }

let tryFindOrigin (tryDecode: ITimelineEvent<EventBody> -> 'event voption, isOrigin: 'event -> bool) xs =

let stack = ResizeArray()

let isOrigin' (u: ITimelineEvent<EventBody>) =

match tryDecode u with

| ValueNone -> false

| ValueSome e ->

stack.Insert(0, e) // WalkResult always renders events ordered correctly - here we're aiming to align with Enum.EventsAndUnfolds

isOrigin e

xs |> Seq.tryFindBack isOrigin', stack.ToArray()

let tryHydrate (tryDecode: ITimelineEvent<EventBody> -> 'event voption, isOrigin: 'event -> bool) (unfolds: Unfold[], etag: string voption) =

match Enum.Unfolds unfolds |> tryFindOrigin (tryDecode, isOrigin) with

| Some u, events ->

let pos = match etag with ValueSome etag -> Position.fromEtagAndIndex (etag, u.Index) | ValueNone -> Position.readOnly

ValueSome (pos, events)

| None, _ -> ValueNone

let feedIteratorMapTi (map: int -> StopwatchInterval -> FeedResponse<'t> -> 'u) (query: FeedIterator<'t>) ct: IAsyncEnumerable<'u> = taskSeq {

// earlier versions, such as 3.9.0, do not implement IDisposable; see linked issue for detail on when SDK team added it

use _ = query // see https://github.com/jet/equinox/issues/225 - in the Cosmos V4 SDK, all this is managed IAsyncEnumerable

let mutable i = 0

while query.HasMoreResults do

let! t, (res: FeedResponse<'t>) = (fun ct -> query.ReadNextAsync(ct)) |> Stopwatch.time ct

yield map i t res

i <- i + 1 }

let private mkQuery (log: ILogger) (container: Container, stream: string) includeTip (maxItems: int) (direction: Direction, minIndex, maxIndex): FeedIterator<Batch> =

let order = if direction = Direction.Forward then "ASC" else "DESC"

let query =

let args = [

match minIndex with None -> () | Some x -> yield "c.n > @minPos", fun (q: QueryDefinition) -> q.WithParameter("@minPos", x)

match maxIndex with None -> () | Some x -> yield "c.i < @maxPos", fun (q: QueryDefinition) -> q.WithParameter("@maxPos", x) ]

let whereClause =

let notTip = sprintf "c.id!=\"%s\"" Tip.WellKnownDocumentId

let conditions = Seq.map fst args

if List.isEmpty args && includeTip then null

else "WHERE " + String.Join(" AND ", if includeTip then conditions else Seq.append conditions (Seq.singleton notTip))

let queryString = $"SELECT c.id, c.i, c._etag, c.n, c.e FROM c %s{whereClause} ORDER BY c.i %s{order}"

let prams = Seq.map snd args

(QueryDefinition queryString, prams) ||> Seq.fold (fun q wp -> q |> wp)

log.Debug("EqxCosmos Query {stream} {query}; n>{minIndex} i<{maxIndex}", stream, query.QueryText, Option.toNullable minIndex, Option.toNullable maxIndex)

container.GetItemQueryIterator<Batch>(query, requestOptions = QueryRequestOptions(PartitionKey = PartitionKey stream, MaxItemCount = maxItems))

// Unrolls the Batches in a response

// NOTE when reading backwards, the events are emitted in reverse Index order to suit the takeWhile consumption

let private mapPage direction (container: Container, streamName: string) (minIndex, maxIndex) (maxRequests: int option)

(log: ILogger) i t (res: FeedResponse<Batch>)

: ITimelineEvent<EventBody>[] * Position option * float =

let log = log |> Log.prop "batchIndex" i

match maxRequests with

| Some mr when i >= mr -> log.Information "batch Limit exceeded"; invalidOp "batch Limit exceeded"

| _ -> ()

let batches, ru = Array.ofSeq res, res.RequestCharge

let unwrapBatch (b: Batch) =

Enum.Events(b, ?minIndex = minIndex, ?maxIndex = maxIndex)

|> if direction = Direction.Backward then Seq.rev else id

let events = batches |> Seq.collect unwrapBatch |> Array.ofSeq

let verbose = log.IsEnabled Events.LogEventLevel.Debug

let count, bytes = events.Length, if verbose then events |> Log.batchLen else 0

let reqMetric: Log.Measurement = { database = container.Database.Id; container = container.Id; stream = streamName; interval = t; bytes = bytes; count = count; ru = ru }

let log = let evt = Log.Metric.QueryResponse (direction, reqMetric) in log |> Log.event evt

let log = if verbose then log |> Log.propEvents events else log

let index = if count = 0 then Nullable () else Nullable <| Seq.min (seq { for x in batches -> x.i })

(log|> Log.prop "bytes" bytes

|> match minIndex with None -> id | Some i -> Log.prop "minIndex" i

|> match maxIndex with None -> id | Some i -> Log.prop "maxIndex" i)

.Information("EqxCosmos {action:l} {count}/{batches} {direction} {ms:f1}ms i={index} {ru}RU",

"Response", count, batches.Length, direction, t.ElapsedMilliseconds, index, ru)

let maybePosition = batches |> Array.tryPick Position.tryFromBatch

events, maybePosition, ru

let private logQuery direction (container: Container, streamName) interval (responsesCount, events: ITimelineEvent<EventBody>[]) n (ru: float) (log: ILogger) =

let verbose = log.IsEnabled Events.LogEventLevel.Debug

let count, bytes = events.Length, if verbose then events |> Log.batchLen else 0

let reqMetric: Log.Measurement = { database = container.Database.Id; container = container.Id; stream = streamName; interval = interval; bytes = bytes; count = count; ru = ru }

let evt = Log.Metric.Query (direction, responsesCount, reqMetric)

let action = match direction with Direction.Forward -> "QueryF" | Direction.Backward -> "QueryB"

(log |> Log.prop "bytes" bytes |> Log.event evt).Information(

"EqxCosmos {action:l} {stream} v{n} {count}/{responses} {ms:f1}ms {ru}RU",

action, streamName, n, count, responsesCount, interval.ElapsedMilliseconds, ru)

let private calculateUsedVersusDroppedPayload stopIndex (xs: ITimelineEvent<EventBody>[]): int * int =

let mutable used, dropped = 0, 0

let mutable found = false

for x in xs do

let bytes = Log.eventLen x

if found then dropped <- dropped + bytes

else used <- used + bytes

if x.Index = stopIndex then found <- true

used, dropped

[<RequireQualifiedAccess; NoComparison; NoEquality>]

type ScanResult<'event> = { found: bool; minIndex: int64; next: int64; maybeTipPos: Position option; events: 'event[] }

let scanTip (tryDecode, isOrigin) (pos: Position, i: int64, xs: #ITimelineEvent<EventBody>[]): ScanResult<'event> =

let origin, events = xs |> Tip.tryFindOrigin (tryDecode, isOrigin)

{ found = Option.isSome origin; maybeTipPos = Some pos; minIndex = i; next = pos.index + 1L; events = events }

// Yields events in ascending Index order

let scan<'event> (log: ILogger) (container, stream) includeTip (maxItems: int) maxRequests direction

(tryDecode: ITimelineEvent<EventBody> -> 'event voption, isOrigin: 'event -> bool)

(minIndex, maxIndex, ct)

: Task<ScanResult<'event> option> = task {

let mutable found = false

let mutable responseCount = 0

let mergeBatches (log: ILogger) (batchesBackward: IAsyncEnumerable<ITimelineEvent<EventBody>[] * Position option * float>) = task {

let mutable lastResponse, maybeTipPos, ru = None, None, 0.

let! events =

batchesBackward

|> TaskSeq.collectSeq (fun (events, maybePos, r) ->

if Option.isNone maybeTipPos then maybeTipPos <- maybePos

lastResponse <- Some events; ru <- ru + r

responseCount <- responseCount + 1

seq { for x in events -> struct (x, tryDecode x) })

|> TaskSeq.takeWhileInclusive (function

| struct (x, ValueSome e) when isOrigin e ->

found <- true

match lastResponse with

| None -> log.Information("EqxCosmos Stop stream={stream} at={index} {case}", stream, x.Index, x.EventType)

| Some batch ->

let used, residual = batch |> calculateUsedVersusDroppedPayload x.Index

log.Information("EqxCosmos Stop stream={stream} at={index} {case} used={used} residual={residual}",

stream, x.Index, x.EventType, used, residual)

false

| _ -> true)

|> TaskSeq.toArrayAsync

return events, maybeTipPos, ru }

let log = log |> Log.prop "batchSize" maxItems |> Log.prop "stream" stream

let readLog = log |> Log.prop "direction" direction

let batches ct: IAsyncEnumerable<ITimelineEvent<EventBody>[] * Position option * float> =

let query = mkQuery readLog (container, stream) includeTip maxItems (direction, minIndex, maxIndex)

feedIteratorMapTi (mapPage direction (container, stream) (minIndex, maxIndex) maxRequests readLog) query ct

let! t, (events, maybeTipPos, ru) = (fun ct -> batches ct |> mergeBatches log) |> Stopwatch.time ct

let raws = Array.map ValueTuple.fst events

let decoded = if direction = Direction.Forward then Array.chooseV ValueTuple.snd events else let xs = Array.chooseV ValueTuple.snd events in Array.Reverse xs; xs

let minMax = (None, raws) ||> Array.fold (fun acc x -> let i = x.Index in Some (match acc with None -> i, i | Some (n, x) -> min n i, max x i))

let version =

match maybeTipPos, minMax with

| Some { index = max }, _

| _, Some (_, max) -> max + 1L

| None, None -> 0L

log |> logQuery direction (container, stream) t (responseCount, raws) version ru