save code

change_proxy/change_cloth_proxy.py

@@ -4,6 +4,10 @@ from flask import Flask, request, jsonify, send_from_directory
 from flask_cors import CORS
 import logging
 from datetime import datetime
+import threading
+import queue
+import uuid
+import time
 
 APP_ROOT = os.path.dirname(os.path.abspath(__file__))
 STATIC_FOLDER = os.path.join(APP_ROOT, 'static')
@@ -25,6 +29,118 @@ def log_message(msg):
     logger.info(msg)
 
 
+class DualQueueProcessor:
+    def __init__(self):
+        self.normal_lock = threading.Lock()
+        self.kuzi_lock = threading.Lock()
+        self.normal_queue = queue.Queue()
+        self.kuzi_queue = queue.Queue()
+        self.normal_processing = 0
+        self.kuzi_processing = 0
+        self.normal_total = 0
+        self.kuzi_total = 0
+        self.stats_lock = threading.Lock()
+        self._start_workers()
+
+    def _start_workers(self):
+        normal_worker = threading.Thread(target=self._process_normal_queue, daemon=True)
+        normal_worker.start()
+        
+        kuzi_worker = threading.Thread(target=self._process_kuzi_queue, daemon=True)
+        kuzi_worker.start()
+        
+        log_message("Dual queue workers started: normal_queue and kuzi_queue")
+
+    def _process_normal_queue(self):
+        while True:
+            task_id, data, result_event, result_container = self.normal_queue.get()
+            log_message(f"[Normal Queue] Processing task {task_id}, queue_size={self.normal_queue.qsize()}")
+            
+            try:
+                response = requests.post(
+                    f"{CHANGE_APP_BASE_URL}/change_cloth_base64",
+                    json=data,
+                    timeout=600
+                )
+                result_container['status_code'] = response.status_code
+                result_container['json'] = response.json()
+                result_container['success'] = True
+                log_message(f"[Normal Queue] Task {task_id} completed with status {response.status_code}")
+            except requests.exceptions.Timeout:
+                result_container['success'] = False
+                result_container['error'] = "Request timeout"
+                log_message(f"[Normal Queue] Task {task_id} timeout")
+            except requests.exceptions.RequestException as e:
+                result_container['success'] = False
+                result_container['error'] = str(e)
+                log_message(f"[Normal Queue] Task {task_id} error: {str(e)}")
+            finally:
+                result_event.set()
+                self.normal_queue.task_done()
+
+    def _process_kuzi_queue(self):
+        while True:
+            task_id, data, result_event, result_container = self.kuzi_queue.get()
+            log_message(f"[Kuzi Queue] Processing task {task_id}, queue_size={self.kuzi_queue.qsize()}")
+            
+            try:
+                response = requests.post(
+                    f"{CHANGE_APP_BASE_URL}/change_cloth_base64",
+                    json=data,
+                    timeout=600
+                )
+                result_container['status_code'] = response.status_code
+                result_container['json'] = response.json()
+                result_container['success'] = True
+                log_message(f"[Kuzi Queue] Task {task_id} completed with status {response.status_code}")
+            except requests.exceptions.Timeout:
+                result_container['success'] = False
+                result_container['error'] = "Request timeout"
+                log_message(f"[Kuzi Queue] Task {task_id} timeout")
+            except requests.exceptions.RequestException as e:
+                result_container['success'] = False
+                result_container['error'] = str(e)
+                log_message(f"[Kuzi Queue] Task {task_id} error: {str(e)}")
+            finally:
+                result_event.set()
+                self.kuzi_queue.task_done()
+
+    def submit_request(self, data):
+        has_kuzi = data.get('kuzi_img') is not None and data.get('kuzi_img') != ''
+        task_id = str(uuid.uuid4())[:8]
+        result_event = threading.Event()
+        result_container = {'success': False}
+        
+        if has_kuzi:
+            with self.stats_lock:
+                self.kuzi_total += 1
+                queue_position = self.kuzi_queue.qsize() + 1
+            self.kuzi_queue.put((task_id, data, result_event, result_container))
+            log_message(f"[Kuzi Queue] Task {task_id} submitted, position={queue_position}")
+            queue_name = "kuzi"
+        else:
+            with self.stats_lock:
+                self.normal_total += 1
+                queue_position = self.normal_queue.qsize() + 1
+            self.normal_queue.put((task_id, data, result_event, result_container))
+            log_message(f"[Normal Queue] Task {task_id} submitted, position={queue_position}")
+            queue_name = "normal"
+        
+        return task_id, result_event, result_container, queue_name, queue_position
+
+    def get_stats(self):
+        with self.stats_lock:
+            return {
+                "normal_queue_size": self.normal_queue.qsize(),
+                "kuzi_queue_size": self.kuzi_queue.qsize(),
+                "normal_total": self.normal_total,
+                "kuzi_total": self.kuzi_total
+            }
+
+
+dual_queue = DualQueueProcessor()
+
+
 @app.route('/')
 def index():
     return send_from_directory(STATIC_FOLDER, 'index.html')
@@ -35,6 +151,16 @@ def health():
     return jsonify({"status": "ok", "service": "change_cloth_proxy"})
 
 
+@app.route('/queue_stats', methods=['GET'])
+def queue_stats():
+    stats = dual_queue.get_stats()
+    return jsonify({
+        "ret": 0,
+        "msg": "success",
+        "data": stats
+    })
+
+
 @app.route('/change_cloth', methods=['POST'])
 def change_cloth():
     log_message("change_cloth called")
@@ -77,22 +203,24 @@ def change_cloth_base64():
         if field not in data:
             return jsonify({"ret": -1, "state": -1, "msg": f"Missing required field: {field}"}), 400
 
-    log_message("Forwarding change_cloth_base64 request")
-
-    try:
-        response = requests.post(
-            f"{CHANGE_APP_BASE_URL}/change_cloth_base64",
-            json=data,
-            timeout=600
-        )
-        log_message(f"change_cloth_base64 response status: {response.status_code}")
-        return jsonify(response.json()), response.status_code
-    except requests.exceptions.Timeout:
-        log_message("change_cloth_base64 request timeout")
-        return jsonify({"ret": -1, "state": -1, "msg": "Request timeout"}), 504
-    except requests.exceptions.RequestException as e:
-        log_message(f"change_cloth_base64 request error: {str(e)}")
-        return jsonify({"ret": -1, "state": -1, "msg": f"Request failed: {str(e)}"}), 500
+    has_kuzi = data.get('kuzi_img') is not None and data.get('kuzi_img') != ''
+    queue_type = "kuzi" if has_kuzi else "normal"
+    
+    task_id, result_event, result_container, queue_name, queue_position = dual_queue.submit_request(data)
+    
+    log_message(f"Task {task_id} queued in {queue_name} queue, position={queue_position}, waiting for result...")
+    
+    if result_event.wait(timeout=610):
+        if result_container.get('success'):
+            log_message(f"Task {task_id} returned successfully")
+            return jsonify(result_container['json']), result_container['status_code']
+        else:
+            error_msg = result_container.get('error', 'Unknown error')
+            log_message(f"Task {task_id} failed: {error_msg}")
+            return jsonify({"ret": -1, "state": -1, "msg": error_msg}), 500
+    else:
+        log_message(f"Task {task_id} timed out waiting in queue")
+        return jsonify({"ret": -1, "state": -1, "msg": "Queue timeout"}), 504
 
 
 @app.route('/do_change_cloth', methods=['POST'])
@@ -127,4 +255,4 @@ def do_change_cloth():
 
 if __name__ == '__main__':
     log_message(f"Starting change_cloth_proxy service on port 5000, forwarding to {CHANGE_APP_BASE_URL}")
-    app.run(host="0.0.0.0", port=5000, debug=True)
+    app.run(host="0.0.0.0", port=5000, debug=True, threaded=True)