From 15c1fc654dacf866b33c03ad5087d5d2cc787367 Mon Sep 17 00:00:00 2001
From: songtianlun <tianlun.song@foxmail.com>
Date: Thu, 13 Feb 2025 16:18:40 +0800
Subject: [PATCH] feat: refine city selection process for weather art

- Implement daily generation limit for weather arts
- Improve city selection by prioritizing active cities
- Introduce logging for city generation start time and list

This update enhances the efficiency of the `BatchGenerateWeatherArtsWorker`. It ensures that the system respects a daily limit on the number of generated images and improves how cities are selected for processing by prioritizing active cities while also conditionally selecting others if slots remain. Additionally, the added logging provides better insights into the generation process.
---
 .../batch_generate_weather_arts_worker.rb     | 93 ++++++++++++++-----
 1 file changed, 70 insertions(+), 23 deletions(-)

diff --git a/app/workers/batch_generate_weather_arts_worker.rb b/app/workers/batch_generate_weather_arts_worker.rb
index a7deda0..79b26dd 100644
--- a/app/workers/batch_generate_weather_arts_worker.rb
+++ b/app/workers/batch_generate_weather_arts_worker.rb
@@ -1,38 +1,85 @@
 class BatchGenerateWeatherArtsWorker
   include Sidekiq::Worker
+
   GENERATION_INTERVAL = 36.hours
   MAX_DURATION = 50.minutes
   SLEEP_DURATION = 120.seconds
-  BATCH_SIZE = 20
-
-  module ProcessType
-    DO_NOTHING = 0
-    GET_WEATHER_WITH_GENERATE_AI_ART = 1
-    GET_WEATHER_ONLY = 2
-  end
+  DAILY_GENERATION_LIMIT = 20  # 每日生成图片上限
 
   def perform(*args)
     start_time = Time.current
-    cities_to_process = get_eligible_cities.shuffle.take(BATCH_SIZE)
-    cities_to_process.each do |city|
+    remaining_slots = calculate_remaining_slots
+    return if remaining_slots <= 0
+
+    # 获取符合条件的城市并处理
+    cities_to_process = select_cities(remaining_slots)&.shuffle
+    print_cities_list(cities_to_process, start_time)
+    process_cities(cities_to_process, start_time)
+  end
+
+  private
+
+  def calculate_remaining_slots
+    today_generations = WeatherArt
+                          .where("DATE(created_at) = ?", Date.today)
+                          .where.not(image_attachment: nil)
+                          .count
+
+    [ DAILY_GENERATION_LIMIT - today_generations, 0 ].max
+  end
+
+  def print_cities_list(cities, start_time)
+    Rails.logger.info "Generate city task start at: #{start_time}"
+    Rails.logger.info "Generate city list: "
+    Rails.logger.info "======================================"
+    Rails.logger.info "\tID\tName\tState\tCountry\tRegion"
+    cities.each do |city|
+      Rails.logger.info "\t#{city.id}\t#{city&.name}\t#{city&.state&.name}\t#{city&.country&.name}\t#{city&.country&.region&.name}"
+    end
+  end
+
+  def select_cities(limit)
+    cutoff_time = Time.current - GENERATION_INTERVAL
+
+    # 基础查询：排除最近生成过的城市
+    base_query = City
+                   .joins("LEFT JOIN (
+        SELECT city_id, MAX(created_at) as last_generation_time
+        FROM weather_arts
+        GROUP BY city_id
+      ) latest_arts ON cities.id = latest_arts.city_id")
+                   .where("latest_arts.last_generation_time IS NULL OR latest_arts.last_generation_time < ?", cutoff_time)
+
+    # 优先选择活跃城市
+    active_cities = base_query
+                      .where(active: true)
+                      .order(:priority)
+                      .limit(limit)
+                      .to_a
+
+    remaining_slots = limit - active_cities.size
+
+    if remaining_slots > 0
+      # 如果还有剩余名额，从其他城市中随机选择
+      other_cities = base_query
+                       .where.not(id: active_cities.map(&:id))
+                       .order("RANDOM()")
+                       .limit(remaining_slots)
+                       .to_a
+
+      active_cities + other_cities
+    else
+      active_cities
+    end
+  end
+
+  def process_cities(cities, start_time)
+    cities.each do |city|
       break if Time.current - start_time > MAX_DURATION
+
       Rails.logger.info "Generating weather art for #{city.name}"
       GenerateWeatherArtWorker.perform_async(city.id)
       sleep SLEEP_DURATION
     end
   end
-
-  private
-
-  def get_eligible_cities
-    cutoff_time = Time.current - GENERATION_INTERVAL
-    City.active
-        .joins("LEFT JOIN (
-          SELECT city_id, MAX(created_at) as last_generation_time
-          FROM weather_arts
-          GROUP BY city_id
-        ) latest_arts ON cities.id = latest_arts.city_id")
-        .where("latest_arts.last_generation_time IS NULL OR latest_arts.last_generation_time < ?", cutoff_time)
-        .order(:priority)
-  end
 end