<commit_before>use crate::{
    background_jobs::Environment,
    models::VersionDownload,
    schema::{crates, metadata, version_downloads, versions},
};

use diesel::prelude::*;
use swirl::PerformError;

#[swirl::background_job]
pub fn update_downloads(env: &Environment) -> Result<(), PerformError> {
    let conn = env.connection()?;
    update(&conn)?;
    Ok(())
}

fn update(conn: &PgConnection) -> QueryResult<()> {
    use self::version_downloads::dsl::*;
    use diesel::dsl::now;
    use diesel::select;

    let rows = version_downloads
        .filter(processed.eq(false))
        .filter(downloads.ne(counted))
        .load(conn)?;

    println!("Updating {} versions", rows.len());
    collect(conn, &rows)?;
    println!("Finished updating versions");

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    diesel::update(version_downloads)
        .set(processed.eq(true))
        .filter(date.lt(diesel::dsl::date(now)))
        .filter(downloads.eq(counted))
        .filter(processed.eq(false))
        .execute(conn)?;
    println!("Finished freezing old version_downloads");

    no_arg_sql_function!(refresh_recent_crate_downloads, ());
    select(refresh_recent_crate_downloads).execute(conn)?;
    println!("Finished running refresh_recent_crate_downloads");

    Ok(())
}

fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
    use diesel::update;

    for download in rows {
        let amt = download.downloads - download.counted;

        conn.transaction::<_, diesel::result::Error, _>(|| {
            // increment the number of counted downloads
            update(version_downloads::table.find(download.id()))
                .set(version_downloads::counted.eq(version_downloads::counted + amt))
                .execute(conn)?;

            // Update the total number of version downloads
            let crate_id = update(versions::table.find(download.version_id))
                .set(versions::downloads.eq(versions::downloads + amt))
                .returning(versions::crate_id)
                .get_result::<i32>(conn)?;

            // Update the total number of crate downloads
            update(crates::table.find(crate_id))
                .set(crates::downloads.eq(crates::downloads + amt))
                .execute(conn)?;

            // Now that everything else for this crate is done, update the global counter of total
            // downloads
            update(metadata::table)
                .set(metadata::total_downloads.eq(metadata::total_downloads + i64::from(amt)))
                .execute(conn)?;

            Ok(())
        })?;
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::{
        env,
        models::{Crate, NewCrate, NewUser, NewVersion, User, Version},
    };
    use std::collections::HashMap;

    fn conn() -> PgConnection {
        let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
        conn.begin_test_transaction().unwrap();
        conn
    }

    fn user(conn: &PgConnection) -> User {
        NewUser::new(2, "login", None, None, "access_token")
            .create_or_update(None, conn)
            .unwrap()
    }

    fn crate_and_version(conn: &PgConnection, user_id: i32) -> (Crate, Version) {<commit_msg>Auto merge of #2154 - jtgeibel:prod/reduce-dl-count-contention, r=sgrif<commit_after>use crate::{
    background_jobs::Environment,
    models::VersionDownload,
    schema::{crates, metadata, version_downloads, versions},
};

use diesel::prelude::*;
use swirl::PerformError;

#[swirl::background_job]
pub fn update_downloads(env: &Environment) -> Result<(), PerformError> {
    let conn = env.connection()?;
    update(&conn)?;
    Ok(())
}

fn update(conn: &PgConnection) -> QueryResult<()> {
    use self::version_downloads::dsl::*;
    use diesel::dsl::now;
    use diesel::select;

    let rows = version_downloads
        .filter(processed.eq(false))
        .filter(downloads.ne(counted))
        .load(conn)?;

    println!("Updating {} versions", rows.len());
    collect(conn, &rows)?;
    println!("Finished updating versions");

    // Anything older than 24 hours ago will be frozen and will not be queried
    // against again.
    diesel::update(version_downloads)
        .set(processed.eq(true))
        .filter(date.lt(diesel::dsl::date(now)))
        .filter(downloads.eq(counted))
        .filter(processed.eq(false))
        .execute(conn)?;
    println!("Finished freezing old version_downloads");

    no_arg_sql_function!(refresh_recent_crate_downloads, ());
    select(refresh_recent_crate_downloads).execute(conn)?;
    println!("Finished running refresh_recent_crate_downloads");

    Ok(())
}

fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
    use diesel::update;

    for download in rows {
        let amt = download.downloads - download.counted;

        conn.transaction::<_, diesel::result::Error, _>(|| {
            // Update the total number of version downloads
            let crate_id = update(versions::table.find(download.version_id))
                .set(versions::downloads.eq(versions::downloads + amt))
                .returning(versions::crate_id)
                .get_result::<i32>(conn)?;

            // Update the total number of crate downloads
            update(crates::table.find(crate_id))
                .set(crates::downloads.eq(crates::downloads + amt))
                .execute(conn)?;

            // Update the global counter of total downloads
            update(metadata::table)
                .set(metadata::total_downloads.eq(metadata::total_downloads + i64::from(amt)))
                .execute(conn)?;

            // Record that these downloads have been propagated to the other tables.  This is done
            // last, immediately before the transaction is committed, to minimize lock contention
            // with counting new downloads.
            update(version_downloads::table.find(download.id()))
                .set(version_downloads::counted.eq(version_downloads::counted + amt))
                .execute(conn)?;

            Ok(())
        })?;
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::{
        env,
        models::{Crate, NewCrate, NewUser, NewVersion, User, Version},
    };
    use std::collections::HashMap;

    fn conn() -> PgConnection {
        let conn = PgConnection::establish(&env("TEST_DATABASE_URL")).unwrap();
        conn.begin_test_transaction().unwrap();
        conn
    }

    fn user(conn: &PgConnection) -> User {
        NewUser::new(2, "login", None, None, "access_token")
            .create_or_update(None, conn)
            .unwrap()
    }

    fn crate_and_version(conn: &PgConnection, user_id: i32) -> (Crate, Version) {
{'fraction_non_alphanumeric': 0.14164471758852795, 'fraction_numerical': 0.003219198127011999, 'mean_word_length': 2.7637665198237884, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 65, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '791698', 'n_tokens_mistral': 2071, 'n_tokens_neox': 2026, 'n_words': 415}