forked from erp-dev/erp
155 lines
4.5 KiB
Python
155 lines
4.5 KiB
Python
"""
|
||
简化的SSE实现，专门解决uvicorn无法退出的问题
|
||
"""
|
||
import json
|
||
import time
|
||
import threading
|
||
from collections import defaultdict
|
||
from django.http import HttpResponse
|
||
from django.views.decorators.csrf import csrf_exempt
|
||
|
||
|
||
# 全局状态
|
||
_clients = {} # 存储客户端连接的状态
|
||
_client_counter = 0
|
||
_shutdown_flag = threading.Event()
|
||
_lock = threading.Lock()
|
||
|
||
|
||
def get_next_client_id():
|
||
"""获取下一个客户端ID"""
|
||
global _client_counter
|
||
with _lock:
|
||
_client_counter += 1
|
||
return _client_counter
|
||
|
||
|
||
def add_client(client_id):
|
||
"""添加客户端"""
|
||
with _lock:
|
||
_clients[client_id] = {
|
||
'active': True,
|
||
'messages': [],
|
||
'last_heartbeat': time.time()
|
||
}
|
||
|
||
|
||
def remove_client(client_id):
|
||
"""移除客户端"""
|
||
with _lock:
|
||
_clients.pop(client_id, None)
|
||
|
||
|
||
def broadcast_message(message):
|
||
"""广播消息给所有客户端"""
|
||
with _lock:
|
||
for client_id, client_data in _clients.items():
|
||
if client_data['active']:
|
||
client_data['messages'].append(message)
|
||
|
||
|
||
def get_client_messages(client_id):
|
||
"""获取客户端的消息"""
|
||
with _lock:
|
||
if client_id in _clients:
|
||
messages = _clients[client_id]['messages'][:]
|
||
_clients[client_id]['messages'].clear()
|
||
_clients[client_id]['last_heartbeat'] = time.time()
|
||
return messages
|
||
return []
|
||
|
||
|
||
def cleanup_clients():
|
||
"""清理所有客户端"""
|
||
global _shutdown_flag
|
||
_shutdown_flag.set()
|
||
with _lock:
|
||
for client_data in _clients.values():
|
||
client_data['active'] = False
|
||
_clients.clear()
|
||
|
||
|
||
@csrf_exempt
|
||
def simple_sse_view(request):
|
||
"""简化的SSE视图"""
|
||
if request.method == 'OPTIONS':
|
||
response = HttpResponse()
|
||
origin = request.META.get('HTTP_ORIGIN')
|
||
if origin:
|
||
response['Access-Control-Allow-Origin'] = origin
|
||
response['Access-Control-Allow-Methods'] = 'GET, OPTIONS'
|
||
response['Access-Control-Allow-Headers'] = 'authorization, content-type, cache-control, accept'
|
||
response['Access-Control-Allow-Credentials'] = 'true'
|
||
response['Access-Control-Max-Age'] = '86400'
|
||
return response
|
||
|
||
def event_generator():
|
||
client_id = get_next_client_id()
|
||
add_client(client_id)
|
||
|
||
try:
|
||
# 发送连接成功消息
|
||
yield f"data: {json.dumps({'type': 'connected', 'client_id': client_id})}\n\n"
|
||
|
||
# 主循环 - 使用更短的检查间隔
|
||
iterations = 0
|
||
while not _shutdown_flag.is_set():
|
||
# 每5次迭代检查一次消息（约0.5秒）
|
||
if iterations % 5 == 0:
|
||
messages = get_client_messages(client_id)
|
||
for message in messages:
|
||
yield f"data: {json.dumps(message)}\n\n"
|
||
|
||
# 每50次迭代发送心跳（约5秒）
|
||
if iterations % 50 == 0:
|
||
yield f": heartbeat {time.time()}\n\n"
|
||
|
||
# 短暂睡眠，让出CPU并允许快速响应关闭信号
|
||
time.sleep(0.1)
|
||
iterations += 1
|
||
|
||
# 如果关闭标志被设置，退出循环
|
||
if _shutdown_flag.is_set():
|
||
break
|
||
|
||
# 发送关闭消息
|
||
yield f"data: {json.dumps({'type': 'shutdown', 'message': 'Server shutting down'})}\n\n"
|
||
|
||
except Exception as e:
|
||
print(f"SSE error for client {client_id}: {e}")
|
||
finally:
|
||
remove_client(client_id)
|
||
print(f"Client {client_id} disconnected")
|
||
|
||
response = HttpResponse(event_generator(), content_type='text/event-stream')
|
||
response['Cache-Control'] = 'no-cache'
|
||
response['X-Accel-Buffering'] = 'no'
|
||
|
||
# CORS headers
|
||
origin = request.META.get('HTTP_ORIGIN')
|
||
if origin:
|
||
response['Access-Control-Allow-Origin'] = origin
|
||
response['Access-Control-Allow-Credentials'] = 'true'
|
||
|
||
return response
|
||
|
||
|
||
def broadcast_test_message():
|
||
"""广播测试消息"""
|
||
message = {
|
||
'type': 'test',
|
||
'message': 'Test message',
|
||
'timestamp': time.time()
|
||
}
|
||
broadcast_message(message)
|
||
return len(_clients)
|
||
|
||
|
||
def get_status():
|
||
"""获取状态"""
|
||
with _lock:
|
||
return {
|
||
'clients': len(_clients),
|
||
'shutdown': _shutdown_flag.is_set()
|
||
}
|