// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

use std::cmp;
use std::collections::btree_map::Entry;
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};

use futures::future::join_all;
use futures::FutureExt;
use futures::{
    future::{self, try_join, try_join3, try_join_all, BoxFuture},
    TryFutureExt,
};
use mz_ore::metrics::MetricsFutureExt;
use mz_ore::{cast::CastFrom, collections::CollectionExt};
use mz_stash_types::{InternalStashError, StashError};
use timely::progress::Antichain;
use timely::PartialOrder;
use tokio::sync::{mpsc, oneshot};
use tokio_postgres::{types::ToSql, Client};

use crate::postgres::{ConsolidateRequest, CountedStatements};
use crate::{AntichainFormatter, AppendBatch, Data, Diff, Id, Stash, StashCollection, Timestamp};

/// The limit AFTER which to split an update batch (that is, we will ship an update that
/// exceeds this number, but then start another batch).
///
/// Cockroach's default limit for sql.conn.max_read_buffer_message_size is 16MiB
/// <https://github.com/cockroachdb/cockroach/blob/7e4e0b195cd61da6cd7a719a5b9aa2e84f68d475/pkg/sql/pgwire/pgwirebase/encoding.go#L50>.
/// Use a number well under that but still big ish that most things won't ever need to
/// batch. Because we are only estimating the value size and ignoring various other
/// things that contribute to the total pgwire message size, having a 14MiB headspace
/// seems safe here.
pub const INSERT_BATCH_SPLIT_SIZE: usize = 2 * 1024 * 1024;

/// [`tokio_postgres`] has a maximum number of arguments it supports when executing a query. This
/// is the limit at which to split a batch to make sure we don't try to include too many elements
/// in any one update.
pub(crate) const MAX_INSERT_ARGUMENTS: u16 = u16::MAX / 4;

impl Stash {
    /// Transactionally executes closure `f`.
    pub async fn with_transaction<F, T>(&mut self, f: F) -> Result<T, StashError>
    where
        F: FnOnce(Transaction) -> BoxFuture<Result<T, StashError>> + Clone + Sync + Send + 'static,
    {
        let (res, mut cons_rx, txn_collections) = self
            .transact(|stmts, client, collections| {
                let f = f.clone();
                let (cons_tx, cons_rx) = mpsc::unbounded_channel();
                let txn_collections = Arc::new(Mutex::new(BTreeMap::new()));
                let tx = Transaction {
                    stmts,
                    client,
                    consolidations: cons_tx,
                    savepoint: Arc::new(Mutex::new(false)),
                    sinces: Arc::new(Mutex::new(BTreeMap::new())),
                    uppers: Arc::new(Mutex::new(BTreeMap::new())),
                    stash_collections: collections,
                    txn_collections: Arc::clone(&txn_collections),
                };
                Box::pin(async move {
                    let res = f(tx).await?;
                    Ok((res, cons_rx, txn_collections))
                })
            })
            .await?;
        while let Some(cons) = cons_rx.recv().await {
            self.sinces_tx
                .send(cons)
                .expect("consolidator unexpectedly gone");
        }
        self.collections
            .extend(std::mem::take(&mut *txn_collections.lock().unwrap()));
        Ok(res)
    }
}

pub struct Transaction<'a> {
    stmts: &'a CountedStatements<'a>,
    client: &'a Client,
    // The set of consolidations that need to be performed if the transaction
    // succeeds.
    consolidations: mpsc::UnboundedSender<ConsolidateRequest>,
    // Savepoint state to enforce the invariant that only one SAVEPOINT is
    // active at once.
    savepoint: Arc<Mutex<bool>>,

    // Cached sinces and uppers for this transaction. These are set on first
    // query and updated on seal/compact.
    sinces: Arc<Mutex<BTreeMap<Id, Antichain<Timestamp>>>>,
    uppers: Arc<Mutex<BTreeMap<Id, Antichain<Timestamp>>>>,

    // Collections cached by the outer Stash.
    stash_collections: &'a BTreeMap<String, Id>,
    // Collections discovered by this transaction.
    txn_collections: Arc<Mutex<BTreeMap<String, Id>>>,
}

impl<'a> Transaction<'a> {
    /// Executes f in a SAVEPOINT. RELEASE if f returns Ok, ROLLBACK if f
    /// returns Err. This must be used for any fn that performs any falliable
    /// operation after its first write. This includes multiple write operations
    /// in a row (any function with multpile writes must use this function).
    async fn in_savepoint<'res, F, T>(&self, f: F) -> Result<T, StashError>
    where
        F: FnOnce() -> BoxFuture<'res, Result<T, StashError>>,
    {
        // Savepoints cannot be used recursively, so panic if that happened.
        {
            // Put this in a block to convince the rust compiler that savepoint
            // won't be used across the await.
            let mut savepoint = self.savepoint.lock().unwrap();
            if *savepoint {
                panic!("cannot call savepoint recursively");
            }
            *savepoint = true;
        };

        self.client.batch_execute("SAVEPOINT txn_savepoint").await?;
        let res = match f().await {
            Ok(t) => {
                self.client
                    .batch_execute("RELEASE SAVEPOINT txn_savepoint")
                    .await?;
                Ok(t)
            }
            Err(err) => {
                self.client
                    .batch_execute("ROLLBACK TO SAVEPOINT txn_savepoint")
                    .await?;
                Err(err)
            }
        };

        let mut savepoint = self.savepoint.lock().unwrap();
        assert!(*savepoint);
        *savepoint = false;

        res
    }

    #[mz_ore::instrument(level = "debug")]
    pub async fn collection<K, V>(&self, name: &str) -> Result<StashCollection<K, V>, StashError>
    where
        K: Data,
        V: Data,
    {
        if let Some(id) = self.stash_collections.get(name) {
            return Ok(StashCollection::new(*id));
        }

        let (statement, histogram) = self.stmts.collection();
        let collection_id_opt: Option<_> = self
            .client
            .query_one(statement, &[&name])
            .wall_time()
            .observe(histogram)
            .await
            .map(|row| row.get("collection_id"))
            .ok();

        let collection_id = match collection_id_opt {
            Some(id) => id,
            None => {
                let collection_id = self
                    .client
                    .query_one(
                        "INSERT INTO collections (name) VALUES ($1) RETURNING collection_id",
                        &[&name],
                    )
                    .await?
                    .get("collection_id");
                self.client
                    .execute(
                        "INSERT INTO sinces (collection_id, since) VALUES ($1, $2)",
                        &[&collection_id, &Timestamp::MIN],
                    )
                    .await?;
                self.client
                    .execute(
                        "INSERT INTO uppers (collection_id, upper) VALUES ($1, $2)",
                        &[&collection_id, &Timestamp::MIN],
                    )
                    .await?;
                collection_id
            }
        };

        self.txn_collections
            .lock()
            .unwrap()
            .insert(name.to_string(), collection_id);
        Ok(StashCollection::new(collection_id))
    }

    /// Returns the ids and names of all collections.
    #[mz_ore::instrument(level = "debug")]
    pub async fn collections(&self) -> Result<BTreeMap<Id, String>, StashError> {
        let rows = self
            .client
            .query("SELECT name, collection_id FROM collections", &[])
            .await?;
        let names = rows
            .into_iter()
            .map(|row| (row.get("collection_id"), row.get("name")));
        Ok(BTreeMap::from_iter(names))
    }

    // TODO(jkosh44) Only used in tests.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn consolidate(&self, id: Id) -> Result<(), StashError> {
        let since = self.since(id).await?;
        self.consolidations
            .send(ConsolidateRequest {
                id,
                since,
                done: None,
            })
            .unwrap();
        Ok(())
    }

    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn upper(
        &self,
        collection_id: Id,
    ) -> Result<Antichain<Timestamp>, StashError> {
        // We can't use .entry here because that would require holding the
        // MutexGuard across the .await.
        if let Some(entry) = self.uppers.lock().unwrap().get(&collection_id) {
            return Ok(entry.clone());
        }
        let (statement, histogram) = self.stmts.upper();
        let upper: Option<Timestamp> = self
            .client
            .query_one(statement, &[&collection_id])
            .wall_time()
            .observe(histogram)
            .await?
            .get("upper");
        let upper = Antichain::from_iter(upper);
        maybe_update_antichain(&self.uppers, collection_id, upper.clone());
        Ok(upper)
    }

    pub(crate) async fn since(
        &self,
        collection_id: Id,
    ) -> Result<Antichain<Timestamp>, StashError> {
        // We can't use .entry here because that would require holding the
        // MutexGuard across the .await.
        if let Some(entry) = self.sinces.lock().unwrap().get(&collection_id) {
            return Ok(entry.clone());
        }
        let (statement, histogram) = self.stmts.since();
        let since: Option<Timestamp> = self
            .client
            .query_one(statement, &[&collection_id])
            .wall_time()
            .observe(histogram)
            .await?
            .get("since");
        let since = Antichain::from_iter(since);
        maybe_update_antichain(&self.sinces, collection_id, since.clone());
        Ok(since)
    }

    /// Iterates over a collection.
    #[mz_ore::instrument(level = "debug")]
    pub async fn iter<K, V>(
        &self,
        collection: StashCollection<K, V>,
    ) -> Result<Vec<((K, V), Timestamp, Diff)>, StashError>
    where
        K: Data,
        V: Data,
    {
        let mut rows = self.iter_raw(collection.id).await?.collect();
        differential_dataflow::consolidation::consolidate_updates(&mut rows);

        // Deserialize the rows into the actual keys and values.
        let rows = rows
            .into_iter()
            .map(|((key, val), ts, diff)| {
                let key = K::decode(key.as_slice()).expect("serialization to roundtrip");
                let val = V::decode(val.as_slice()).expect("serialization to roundtrip");

                ((key, val), ts, diff)
            })
            .collect();

        Ok(rows)
    }

    /// Iterates over a collection, returning the raw data on disk, unconsolidated.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn iter_raw(
        &self,
        id: Id,
    ) -> Result<impl Iterator<Item = ((Vec<u8>, Vec<u8>), Timestamp, Diff)>, StashError> {
        let since = match self.since(id).await?.into_option() {
            Some(since) => since,
            None => {
                return Err(StashError::from(
                    "cannot iterate collection with empty since frontier",
                ));
            }
        };
        let (statement, histogram) = self.stmts.iter();
        let rows = self
            .client
            .query(statement, &[&id])
            .wall_time()
            .observe(histogram)
            .await?
            .into_iter()
            .map(move |row| {
                let key = row.get("key");
                let value = row.get("value");
                let time = row.get("time");
                let diff: Diff = row.get("diff");
                ((key, value), cmp::max(time, since), diff)
            });
        Ok(rows)
    }

    /// Iterates over the values of a key.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn iter_key<K, V>(
        &self,
        collection: StashCollection<K, V>,
        key: &K,
    ) -> Result<Vec<(V, Timestamp, Diff)>, StashError>
    where
        K: Data,
        V: Data,
    {
        let key = key.encode_to_vec();
        let (statement, histogram) = self.stmts.iter_key();
        let (since, rows) = future::try_join(
            self.since(collection.id),
            self.client
                .query(statement, &[&collection.id, &key])
                .map_err(|err| err.into())
                .wall_time()
                .observe(histogram),
        )
        .await?;
        let since = match since.into_option() {
            Some(since) => since,
            None => {
                return Err(StashError::from(
                    "cannot iterate collection with empty since frontier",
                ));
            }
        };
        let mut rows = rows
            .into_iter()
            .map(|row| {
                let value: Vec<u8> = row.try_get("value")?;
                let value = V::decode(value.as_slice())?;
                let time = row.try_get("time")?;
                let diff = row.try_get("diff")?;
                Ok::<_, StashError>((value, cmp::max(time, since), diff))
            })
            .collect::<Result<Vec<_>, _>>()?;
        differential_dataflow::consolidation::consolidate_updates(&mut rows);
        Ok(rows)
    }

    /// Returns the most recent timestamp at which sealed entries can be read.
    #[mz_ore::instrument(level = "debug")]
    async fn peek_timestamp_id(&self, id: Id) -> Result<Timestamp, StashError> {
        let (since, upper) = try_join(self.since(id), self.upper(id)).await?;
        if PartialOrder::less_equal(&upper, &since) {
            return Err(StashError {
                inner: InternalStashError::PeekSinceUpper(format!(
                    "collection {} since {} is not less than upper {}",
                    id,
                    AntichainFormatter(&since),
                    AntichainFormatter(&upper)
                )),
            });
        }
        match upper.as_option() {
            Some(ts) => match ts.checked_sub(1) {
                Some(ts) => Ok(ts),
                None => Err("could not determine peek timestamp".into()),
            },
            None => Ok(Timestamp::MAX),
        }
    }

    /// Returns the most recent timestamp at which sealed entries can be read.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn peek_timestamp<K, V>(
        &self,
        collection: StashCollection<K, V>,
    ) -> Result<Timestamp, StashError>
    where
        K: Data,
        V: Data,
    {
        self.peek_timestamp_id(collection.id).await
    }

    /// Returns the current value of sealed entries.
    ///
    /// Entries are iterated in `(key, value)` order and are guaranteed to be
    /// consolidated.
    ///
    /// Sealed entries are those with timestamps less than the collection's upper
    /// frontier.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn peek<K, V>(
        &self,
        collection: StashCollection<K, V>,
    ) -> Result<Vec<(K, V, Diff)>, StashError>
    where
        K: Data,
        V: Data,
    {
        let timestamp = self.peek_timestamp(collection).await?;
        let mut rows: Vec<_> = self
            .iter(collection)
            .await?
            .into_iter()
            .filter_map(|((k, v), data_ts, diff)| {
                if data_ts.less_equal(&timestamp) {
                    Some((k, v, diff))
                } else {
                    None
                }
            })
            .collect();
        differential_dataflow::consolidation::consolidate_updates(&mut rows);
        Ok(rows)
    }

    /// Returns the current k,v pairs of sealed entries, erroring if there is more
    /// than one entry for a given key or the multiplicity is not 1 for each key.
    ///
    /// Sealed entries are those with timestamps less than the collection's upper
    /// frontier.
    #[mz_ore::instrument(level = "debug")]
    pub async fn peek_one<K, V>(
        &self,
        collection: StashCollection<K, V>,
    ) -> Result<BTreeMap<K, V>, StashError>
    where
        K: Data,
        V: Data,
    {
        let rows = self.peek(collection).await?;
        let mut res = BTreeMap::new();
        for (k, v, diff) in rows {
            if diff != 1 {
                return Err(format!("unexpected peek multiplicity of {diff}").into());
            }
            if res.insert(k, v).is_some() {
                return Err(format!("duplicate peek keys for collection {}", collection.id).into());
            }
        }
        Ok(res)
    }

    /// Returns the current sealed value for the given key, erroring if there is
    /// more than one entry for the key or its multiplicity is not 1.
    ///
    /// Sealed entries are those with timestamps less than the collection's upper
    /// frontier.
    #[mz_ore::instrument(level = "debug")]
    pub async fn peek_key_one<K, V>(
        &self,
        collection: StashCollection<K, V>,
        key: &K,
    ) -> Result<Option<V>, StashError>
    where
        K: Data,
        V: Data,
    {
        let timestamp = self.peek_timestamp(collection).await?;
        let mut rows: Vec<_> = self
            .iter_key(collection, key)
            .await?
            .into_iter()
            .filter_map(|(v, data_ts, diff)| {
                if data_ts.less_equal(&timestamp) {
                    Some((v, diff))
                } else {
                    None
                }
            })
            .collect();
        differential_dataflow::consolidation::consolidate(&mut rows);
        let v = match rows.len() {
            1 => {
                let (v, diff) = rows.into_element();
                match diff {
                    1 => Some(v),
                    0 => None,
                    _ => return Err("multiple values unexpected".into()),
                }
            }
            0 => None,
            _ => return Err("multiple values unexpected".into()),
        };
        Ok(v)
    }

    /// Applies batches to the current transaction. If any batch fails and in
    /// error returned, all other applications are rolled back.
    #[mz_ore::instrument(level = "debug")]
    pub async fn append(
        &self,
        batches: Vec<AppendBatch>,
    ) -> Result<BoxFuture<'static, ()>, StashError> {
        if batches.is_empty() {
            return Ok(Box::pin(futures::future::ready(())));
        }

        let consolidations = self
            .in_savepoint(|| {
                Box::pin(async move {
                    let futures = batches.into_iter().map(
                        |AppendBatch {
                             collection_id,
                             lower,
                             upper,
                             entries,
                             ..
                         }| async move {
                            // Clone to appease rust async.
                            let compact = lower.clone();
                            let lower1 = lower.clone();
                            let lower2 = lower.clone();
                            let upper1 = upper.clone();
                            // The new upper must be validated before being sealed.
                            // Compaction can be done in any order because none of
                            // update/seal care about what the since is. Update, although it
                            // cares about the upper, can also be done in any order because
                            // we pass in the old upper, so it's not a problem if the seal
                            // has executed already since update won't query the current
                            // upper.
                            try_join3(
                                    async move {
                                        let current_upper = self.upper(collection_id).await?;
                                        if current_upper != lower1 {
                                            return Err(StashError::from(format!(
                                                "unexpected lower, got {:?}, expected {:?}",
                                                current_upper, lower1
                                            )));
                                        }

                                        self.seal(collection_id, upper1, Some(lower)).await
                                    },
                                    async move {
                                        self.update(collection_id, &entries, Some(lower2)).await
                                    },
                                    async move {
                                        self.compact(collection_id, &compact, Some(upper)).await
                                    },
                                )
                                .await?;

                            let (tx, rx) = oneshot::channel();
                            let request = ConsolidateRequest {
                                id: collection_id,
                                since: self.since(collection_id).await?,
                                done: Some(tx),
                            };

                            Ok::<_, StashError>((request, rx))
                        },
                    );
                    try_join_all(futures).await
                })
            })
            .await?;

        let mut notifs = Vec::with_capacity(consolidations.len());
        for (req, notif) in consolidations {
            self.consolidations.send(req).unwrap();
            notifs.push(notif);
        }

        // Note: we ignore the result of the Consolidation requests because consolidation cannot fail.
        Ok(Box::pin(join_all(notifs).map(|_| ())))
    }

    /// Like update, but starts a savepoint.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn update_savepoint<K, V>(
        &self,
        collection_id: Id,
        entries: &[((K, V), Timestamp, Diff)],
        upper: Option<Antichain<Timestamp>>,
    ) -> Result<(), StashError>
    where
        K: Data,
        V: Data,
    {
        let entries: Vec<_> = entries
            .into_iter()
            .map(|((key, val), time, diff)| {
                let key = key.encode_to_vec();
                let val = val.encode_to_vec();
                ((key, val), *time, *diff)
            })
            .collect();
        self.in_savepoint(|| {
            Box::pin(async { self.update(collection_id, &entries[..], upper).await })
        })
        .await
    }

    /// Directly add k, v, ts, diff tuples to a collection. `upper` can be `Some`
    /// if the collection's upper is already known. Caller must have already
    /// called in_savepoint.
    ///
    /// This function should not be called outside of the stash crate since it
    /// allows for arbitrary bytes, non-unit diffs in collections, and doesn't
    /// support transaction safety. Use `TypedCollection`'s methods instead.
    #[mz_ore::instrument(level = "debug")]
    async fn update(
        &self,
        collection_id: Id,
        entries: &[((Vec<u8>, Vec<u8>), Timestamp, Diff)],
        upper: Option<Antichain<Timestamp>>,
    ) -> Result<(), StashError> {
        {
            // Panic if the caller didn't initiate the savepoint.
            let savepoint = self.savepoint.lock().unwrap();
            assert!(*savepoint);
        }

        if entries.is_empty() {
            return Ok(());
        }

        // Check the upper in a separate future so we can issue the updates without
        // waiting for it first.
        let upper_fut = async {
            let upper = match upper {
                Some(upper) => upper,
                None => self.upper(collection_id).await?,
            };
            for ((_key, _value), time, _diff) in entries {
                if !upper.less_equal(time) {
                    return Err(StashError::from(format!(
                        "entry time {} is less than the current upper frontier {}",
                        time,
                        AntichainFormatter(&upper)
                    )));
                }
            }
            Ok(upper)
        };

        let insert_fut = async {
            let mut entries = entries.iter();
            loop {
                let mut args: Vec<&'_ (dyn ToSql + Sync)> = Vec::new();
                // All rows use the collection id, so hard code it as the first.
                args.push(&collection_id);
                let mut batch_size = 0;
                let mut estimated_json_size = 0;
                // Accumulate into a batch until the size limit is exceeded or there are no more
                // entries.
                while let Some(((key, value), time, diff)) = entries.next() {
                    estimated_json_size += key.len();
                    estimated_json_size += value.len();
                    args.push(key);
                    args.push(value);
                    args.push(time);
                    args.push(diff);
                    batch_size += 1;
                    if estimated_json_size > INSERT_BATCH_SPLIT_SIZE {
                        break;
                    }
                    if args.len() > CastFrom::cast_from(MAX_INSERT_ARGUMENTS) {
                        break;
                    }
                }
                if batch_size == 0 {
                    return Ok(());
                }
                let stmt = self.stmts.update(self.client, batch_size).await?;
                self.client.execute(&stmt, &args).await?;
            }
        };
        try_join(upper_fut, insert_fut).await?;
        Ok(())
    }

    /// Sets the since of a collection. The current upper can be `Some` if it is
    /// already known.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn compact<'ts>(
        &self,
        id: Id,
        new_since: &'ts Antichain<Timestamp>,
        upper: Option<Antichain<Timestamp>>,
    ) -> Result<(), StashError> {
        // Do all validation first.
        try_join(
            async move {
                let since = self.since(id).await?;
                if PartialOrder::less_than(new_since, &since) {
                    return Err(StashError::from(format!(
                        "compact request {} is less than the current since frontier {}",
                        AntichainFormatter(new_since),
                        AntichainFormatter(&since)
                    )));
                }
                Ok(())
            },
            async move {
                let upper = match upper {
                    Some(upper) => upper,
                    None => self.upper(id).await?,
                };
                if PartialOrder::less_than(&upper, new_since) {
                    return Err(StashError::from(format!(
                        "compact request {} is greater than the current upper frontier {}",
                        AntichainFormatter(new_since),
                        AntichainFormatter(&upper)
                    )));
                }
                Ok(())
            },
        )
        .await?;

        // If successful, execute the change in the txn.
        let (statement, histogram) = self.stmts.compact();
        self.client
            .execute(statement, &[&new_since.as_option(), &id])
            .map_err(StashError::from)
            .wall_time()
            .observe(histogram)
            .await?;
        maybe_update_antichain(&self.sinces, id, new_since.clone());
        Ok(())
    }

    /// Sets the upper of a collection. The current upper can be `Some` if it is
    /// already known.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn seal(
        &self,
        id: Id,
        new_upper: Antichain<Timestamp>,
        upper: Option<Antichain<Timestamp>>,
    ) -> Result<(), StashError> {
        let upper = match upper {
            Some(upper) => upper,
            None => self.upper(id).await?,
        };
        if PartialOrder::less_than(&new_upper, &upper) {
            return Err(StashError::from(format!(
                "seal request {} is less than the current upper frontier {}",
                AntichainFormatter(&new_upper),
                AntichainFormatter(&upper),
            )));
        }

        let (statement, histogram) = self.stmts.seal();
        self.client
            .execute(statement, &[&new_upper.as_option(), &id])
            .map_err(StashError::from)
            .wall_time()
            .observe(histogram)
            .await?;
        maybe_update_antichain(&self.uppers, id, new_upper);
        Ok(())
    }
}

/// Updates an antichain cache if the new value has advanced. Needed because the
/// functions here are often called in try_joins where the futures execute in
/// unknown order and we want to prevent a race condition poisioning the cache by
/// going backward.
fn maybe_update_antichain(
    map: &Arc<Mutex<BTreeMap<Id, Antichain<Timestamp>>>>,
    id: Id,
    updated: Antichain<Timestamp>,
) {
    match map.lock().unwrap().entry(id) {
        Entry::Occupied(mut entry) => {
            if PartialOrder::less_than(entry.get(), &updated) {
                entry.insert(updated);
            }
        }
        Entry::Vacant(entry) => {
            entry.insert(updated);
        }
    };
}