Fix download counts (#746) (#747)

* Fix download counts (#746)

* Fix download counts

* remove unsafe send

* update indexing time

* run prep

* run prep again

src/queue/analytics.rs

@@ -1,16 +1,16 @@
+use crate::database::models::DatabaseError;
 use crate::models::analytics::{Download, PageView, Playtime};
-use dashmap::DashSet;
+use crate::routes::ApiError;
+use dashmap::{DashMap, DashSet};
+use redis::cmd;
+use sqlx::PgPool;
+use crate::database::redis::RedisPool;
 
-#[cfg(test)]
-mod tests;
-
-const VIEWS_TABLENAME: &str = "views";
-const DOWNLOADS_TABLENAME: &str = "downloads";
-const PLAYTIME_TABLENAME: &str = "playtime";
+const DOWNLOADS_NAMESPACE: &str = "downloads";
 
 pub struct AnalyticsQueue {
     views_queue: DashSet<PageView>,
-    downloads_queue: DashSet<Download>,
+    downloads_queue: DashMap<String, Download>,
     playtime_queue: DashSet<Playtime>,
 }
 
@@ -25,7 +25,7 @@ impl AnalyticsQueue {
     pub fn new() -> Self {
         AnalyticsQueue {
             views_queue: DashSet::with_capacity(1000),
-            downloads_queue: DashSet::with_capacity(1000),
+            downloads_queue: DashMap::with_capacity(1000),
             playtime_queue: DashSet::with_capacity(1000),
         }
     }
@@ -35,45 +35,138 @@ impl AnalyticsQueue {
     }
 
     pub fn add_download(&self, download: Download) {
-        self.downloads_queue.insert(download);
+        let octets = download.ip.octets();
+        let ip_stripped = u64::from_be_bytes([
+            octets[0], octets[1], octets[2], octets[3], octets[4], octets[5], octets[6], octets[7],
+        ]);
+        self.downloads_queue
+            .insert(format!("{}-{}", ip_stripped, download.project_id), download);
     }
 
     pub fn add_playtime(&self, playtime: Playtime) {
         self.playtime_queue.insert(playtime);
     }
 
-    pub async fn index(&self, client: clickhouse::Client) -> Result<(), clickhouse::error::Error> {
-        Self::index_queue(&client, &self.views_queue, VIEWS_TABLENAME).await?;
-        Self::index_queue(&client, &self.downloads_queue, DOWNLOADS_TABLENAME).await?;
-        Self::index_queue(&client, &self.playtime_queue, PLAYTIME_TABLENAME).await?;
+    pub async fn index(
+        &self,
+        client: clickhouse::Client,
+        redis: &RedisPool,
+        pool: &PgPool,
+    ) -> Result<(), ApiError> {
+        let views_queue = self.views_queue.clone();
+        self.views_queue.clear();
+
+        let downloads_queue = self.downloads_queue.clone();
+        self.downloads_queue.clear();
+
+        let playtime_queue = self.playtime_queue.clone();
+        self.playtime_queue.clear();
+
+        if !views_queue.is_empty() {
+            let mut views = client.insert("views")?;
+
+            for view in views_queue {
+                views.write(&view).await?;
+            }
+
+            views.end().await?;
+        }
+
+        if !playtime_queue.is_empty() {
+            let mut playtimes = client.insert("playtime")?;
+
+            for playtime in playtime_queue {
+                playtimes.write(&playtime).await?;
+            }
+
+            playtimes.end().await?;
+        }
+
+        if !downloads_queue.is_empty() {
+            let mut downloads_keys = Vec::new();
+            let raw_downloads = DashMap::new();
+
+            for (index, (key, download)) in downloads_queue.into_iter().enumerate() {
+                downloads_keys.push(key);
+                raw_downloads.insert(index, download);
+            }
+
+            let mut redis = redis.pool.get().await.map_err(DatabaseError::RedisPool)?;
+
+            let results = cmd("MGET")
+                .arg(
+                    downloads_keys
+                        .iter()
+                        .map(|x| format!("{}:{}", DOWNLOADS_NAMESPACE, x))
+                        .collect::<Vec<_>>(),
+                )
+                .query_async::<_, Vec<Option<u32>>>(&mut redis)
+                .await
+                .map_err(DatabaseError::CacheError)?;
+
+            let mut pipe = redis::pipe();
+            for (idx, count) in results.into_iter().enumerate() {
+                let key = &downloads_keys[idx];
+
+                let new_count = if let Some(count) = count {
+                    if count > 5 {
+                        raw_downloads.remove(&idx);
+                        continue;
+                    }
+
+                    count + 1
+                } else {
+                    1
+                };
+
+                pipe.atomic().set_ex(
+                    format!("{}:{}", DOWNLOADS_NAMESPACE, key),
+                    new_count,
+                    6 * 60 * 60,
+                );
+            }
+            pipe.query_async(&mut *redis)
+                .await
+                .map_err(DatabaseError::CacheError)?;
+
+            let version_ids = raw_downloads
+                .iter()
+                .map(|x| x.version_id as i64)
+                .collect::<Vec<_>>();
+            let project_ids = raw_downloads
+                .iter()
+                .map(|x| x.project_id as i64)
+                .collect::<Vec<_>>();
+
+            let mut transaction = pool.begin().await?;
+            let mut downloads = client.insert("downloads")?;
+
+            for (_, download) in raw_downloads {
+                downloads.write(&download).await?;
+            }
+
+            sqlx::query!(
+                "UPDATE versions
+                SET downloads = downloads + 1
+                WHERE id = ANY($1)",
+                &version_ids
+            )
+                .execute(&mut *transaction)
+                .await?;
+
+            sqlx::query!(
+                "UPDATE mods
+                SET downloads = downloads + 1
+                WHERE id = ANY($1)",
+                &project_ids
+            )
+                .execute(&mut *transaction)
+                .await?;
+
+            transaction.commit().await?;
+            downloads.end().await?;
+        }
 
         Ok(())
     }
-
-    async fn index_queue<T>(
-        client: &clickhouse::Client,
-        queue: &DashSet<T>,
-        table_name: &str,
-    ) -> Result<(), clickhouse::error::Error>
-    where
-        T: serde::Serialize + Eq + std::hash::Hash + Clone + clickhouse::Row,
-    {
-        if queue.is_empty() {
-            return Ok(());
-        }
-
-        let current_queue = queue.clone();
-        queue.clear();
-
-        let mut inserter = client.inserter(table_name)?;
-
-        for row in current_queue {
-            inserter.write(&row).await?;
-            inserter.commit().await?;
-        }
-
-        inserter.end().await?;
-
-        Ok(())
-    }
-}
+}