Websocket 적용기

NextJS, FastAPI Websocket 구현기

커뮤니티 기능을 도입하며 사용자 알림기능이 필요하게 되면서 Websocket을 구현하게 되었습니다.

설계

Kafka 2.13-4.1.1 Kraft Mode

설계는 간단합니다.

FastAPI, Kafka, Redis 에서의 흐름입니다.

이후의 FastAPI, WebSocket 에서의 흐름입니다.

Redis를 껴서 관리하면 편할 것 같아 적용했고, table에 저장하는 이유는 사용자 알림 이력 관리 입니다.

구현 순서는 다음과 같습니다.

1. 테이블 생성
2. Kafka Topic 생성
3. Kafka Consumer 생성
4. FastAPI -> Kafka Produce 로직 구현
5. FastAPI - Redis Listener 및 Websocket 전달 구현
6. NextJS Websocket 구현

Table 생성

알림 정보를 저장하는 테이블 입니다.

별도의 테이블을 조회할 필요 없이, 해당 데이터만 가지고 알림을 바로 보여줄 수 있게 JSON 구조를 사용합니다.

CREATE TABLE IF NOT EXISTS user_notifications (
    id SERIAL PRIMARY KEY,
    user_email TEXT NOT NULL,
    noti_type TEXT NOT NULL,
    payload JSONB,
    is_read BOOLEAN DEFAULT FALSE,
    created_time TIMESTAMP DEFAULT NOW()
);
COMMENT ON COLUMN user_notifications.id IS '자동 생성 ID';
COMMENT ON COLUMN user_notifications.user_email IS '사용자 이메일';
COMMENT ON COLUMN user_notifications.noti_type IS '알림 타입';
COMMENT ON COLUMN user_notifications.payload IS '알림 내용';
COMMENT ON COLUMN user_notifications.is_read IS '알림 읽기 여부';
COMMENT ON COLUMN user_notifications.created_time IS '알림 생성 시간';

Kafka 구현

구현에 들어가겠습니다.

구조

consumer/
- handelers/
-- notification_handlers.py
- config.py
- consumer_base.py
- logger.py
- pg_client.py
main_notification_consumer.py

Kafka Topic 생성

간단하게 토픽을 하나 추가해줍니다.

user-notification 이라는 토픽을 생성했습니다.

/home/kafka/kafka_2.13-4.1.1/bin/kafka-topics.sh \
  --create \
  --topic user-notification \
  --bootstrap-server localhost:9092 \
  --partitions 3 \
  --replication-factor 1 \
  --config retention.ms=604800000

Kafka Consumer 생성

별도의 Python 프로세스로 실행할 Consumer 입니다.

여기에서는 데이터를 받은 후 notification table에 적재한 뒤, redis에 저장합니다.

config.py

.env의 정보를 가져와서 변수로 두는 곳입니다.

import os
from dotenv import load_dotenv

load_dotenv()

LOG_DIR = os.getenv("LOG_DIR", "./logs")
LOG_FILE = os.path.join(LOG_DIR, "consumer.log")

# PostgreSQL
PG_CONFIG = {
    "host": os.getenv("DB_HOST"),
    "port": os.getenv("DB_PORT"),
    "dbname": os.getenv("DB_NAME"),
    "user": os.getenv("DB_USER"),
    "password": os.getenv("DB_PASSWORD"),
}

# Kafka
NOTIFICATION_KAFKA_CONFIG = {
    "bootstrap.servers": os.getenv("BOOTSTRAP_SERVER"),
    "group.id": os.getenv("GROUP_ID"),
    "auto.offset.reset": os.getenv("OFFSET_RESET_CONFIG"),
    "topic": os.getenv("NOTIFICATION_TOPIC"),
}

pg_client.py

DB 연결부 입니다.

import psycopg2
from consumer.config import PG_CONFIG
from consumer.logger import logger


def get_pg_connection():
    return psycopg2.connect(**PG_CONFIG)

consumer_base.py

web log와 notification에서 공통적으로 사용하는 conusmer 뼈대입니다.

import json
from confluent_kafka import Consumer, KafkaError
from consumer.logger import logger


def run_consumer(kafka_config, process_message_callback):
    consumer = Consumer(
        {
            "bootstrap.servers": kafka_config["bootstrap.servers"],
            "group.id": kafka_config["group.id"],
            "auto.offset.reset": kafka_config["auto.offset.reset"],
        }
    )
    topic = kafka_config["topic"]
    consumer.subscribe([topic])
    logger.info(f"Subscribed to topic: {topic}")

    try:
        while True:
            msg = consumer.poll(1.0)
            if msg is None:
                continue
            if msg.error():
                if msg.error().code() == KafkaError._PARTITION_EOF:
                    logger.warning(
                        f"End of partition: {msg.topic()} [{msg.partition()}]"
                    )
                else:
                    logger.error(f"Kafka error: {msg.error().str()}")
                continue

            try:
                data = json.loads(msg.value().decode("utf-8"))
                logger.info(f"Received message JSON: {data}")

                process_message_callback(data)

            except json.JSONDecodeError as e:
                logger.error(f"JSON 디코딩 실패: {e}")
            except Exception as e:
                logger.error(f"메시지 처리 실패: {e}")

    except KeyboardInterrupt:
        logger.info("Consumer 종료됨")
    finally:
        consumer.close()
        logger.info("Consumer 연결 종료")

notification_handler.py

notification handler 입니다.

FastAPI의 요청을 받은 후 요청 타입에 따라 개별 쿼리로 테이블에 저장한 뒤 Redis로 요청을 보내고 있습니다.

요청 타입은 다음과 같습니다.

• 팔로우 한 사람이 게시글 작성한 경우
• 내 게시글에 댓글이 달린 경우
• 내 댓글에 대댓글이 달린 경우
• 사이트에서 제재를 당한 경우
• 누군가 나를 팔로우 한 경우

import json
import redis
from consumer.pg_client import get_pg_connection
from consumer.logger import logger
import os
from dotenv import load_dotenv

load_dotenv()
pg_conn = get_pg_connection()
logger.info("알림 DB 연결 성공")

# Redis 연결
redis_client = redis.StrictRedis(
    host=os.getenv("REDIS_HOST"),
    port=int(os.getenv("REDIS_PORT")),
    db=int(os.getenv("REDIS_DB")),
    decode_responses=True,
)
logger.info("Redis 연결 성공")

NOTIFICATION_HANDLERS = {
    "create_post": {
        "query": """
            INSERT INTO user_notifications (user_email, noti_type, payload)
            SELECT uf.following_email, %(noti_type)s, %(payload)s
            FROM user_follows uf
            WHERE uf.follower_email = %(author_email)s
            RETURNING user_email
        """,
        "param_builder": lambda data: {
            "noti_type": data["noti_type"],
            "payload": json.dumps(data),
            "author_email": data["author_email"],
        },
    },
    "create_parent_comment": {
        "query": """
            INSERT INTO user_notifications (user_email, noti_type, payload)
            SELECT cp.user_email, %(noti_type)s, %(payload)s
            FROM community_posts cp
            WHERE cp.id = %(post_id)s
            RETURNING user_email
        """,
        "param_builder": lambda data: {
            "noti_type": data["noti_type"],
            "payload": json.dumps(data),
            "post_id": data["post_id"],
        },
    },
    "create_child_comment": {
        "query": """
            INSERT INTO user_notifications (user_email, noti_type, payload)
            SELECT cc.user_email, %(noti_type)s, %(payload)s
            FROM community_comments cc
            WHERE cc.id = %(parent_comment_id)s
            RETURNING user_email
        """,
        "param_builder": lambda data: {
            "noti_type": data["noti_type"],
            "payload": json.dumps(data),
            "parent_comment_id": data["parent_comment_id"],
        },
    },
    "follow_user": {
        "query": """
            INSERT INTO user_notifications (user_email, noti_type, payload)
            SELECT ui.email, %(noti_type)s, %(payload)s
            FROM user_info ui
            WHERE ui.email = %(follower_email)s
            RETURNING user_email
        """,
        "param_builder": lambda data: {
            "noti_type": data["noti_type"],
            "payload": json.dumps(data),
            "follower_email": data["follower_email"],
        },
    },
    "penalty_user": {
        "query": """
            INSERT INTO user_notifications (user_email, noti_type, payload)
            VALUES (%(user_email)s, %(noti_type)s, %(payload)s)
            RETURNING user_email
        """,
        "param_builder": lambda data: {
            "user_email": data["user_email"],
            "noti_type": data["noti_type"],
            "payload": json.dumps(data),
        },
    },
}


def save_notifications_and_push_to_redis(cur, query, params, data):
    """
    공통 함수: DB에 알림 저장 후 Redis에 push
    """
    cur.execute(query, params)
    notified_users = [row[0] for row in cur.fetchall()]
    pg_conn.commit()

    for user_email in notified_users:
        redis_key = f"notifications:{user_email}"
        redis_client.lpush(redis_key, json.dumps(data))
        redis_client.expire(redis_key, 60 * 60 * 24 * 7)  # TTL 7일
        # WebSocket용 pub/sub push
        redis_client.publish(f"notifications_channel:{user_email}", json.dumps(data))

    return notified_users


def process_notification_message(data):
    """
    알림 처리하기
    """
    handler = NOTIFICATION_HANDLERS.get(data["noti_type"])
    if not handler:
        logger.warning(f"지원하지 않는 알림 타입: {data['noti_type']}")
        return

    try:
        with pg_conn.cursor() as cur:
            insert_query = handler["query"]
            params = handler["param_builder"](data)
            save_notifications_and_push_to_redis(cur, insert_query, params, data)
        logger.info(f"알림 처리 완료: {data} (Redis Save)")
    except Exception as e:
        logger.error(f"알림 처리 실패: {e}")
        pg_conn.rollback()

main_notification_handler.py

프로세스를 실행하는 main 함수입니다.

from consumer.config import NOTIFICATION_KAFKA_CONFIG
from consumer.consumer_base import run_consumer
from consumer.handlers.notification_handler import process_notification_message

if __name__ == "__main__":
    run_consumer(NOTIFICATION_KAFKA_CONFIG, process_notification_message)

FastAPI 구현

FastAPI에서는 UI 요청별로 Kafka에 Produce를 하고, Redis에 Listen하다가 데이터 변경이 감지되면 Websocket으로 전달하는 역할을 하고 있습니다.

구조

api/
- comment/
-- service.py
util/
- kafka_producer.py
- websocket.py
main.py

kafka_producer.py

메시지를 Kafka로 Produce하는 함수입니다.

from confluent_kafka import Producer
import os
from dotenv import load_dotenv
import logging

load_dotenv()

BOOTSTRAP_SERVER = os.getenv("BOOTSTRAP_SERVER")
NOTIFICATION_TOPIC = os.getenv("NOTIFICATION_TOPIC")

producer_conf = {"bootstrap.servers": BOOTSTRAP_SERVER, "client.id": "fastapi-producer"}
producer = Producer(producer_conf)
logger = logging.getLogger("kafka.producer")


def delivery_report(err, msg):
    if err:
        logger.error(f"Delivery failed: {err}")


def produce_notification(value: str):
    producer.produce(
        NOTIFICATION_TOPIC, value=value.encode("utf-8"), callback=delivery_report
    )
    producer.poll(0)
    producer.flush(0.2)

service에서 produce 적용

기존 사용 함수 중간에 produce 하는 부분을 추가하면 됩니다.

본인이 작성한 것은 무시하게 해두었습니다.

    @staticmethod
    def insert_parent_comment(request_info: InsertParentComment, user_email: str):
        try:
            with DataBaseConnector.SessionLocal() as s:
                # nano id 생성
                new_id = generate(size=12)
                bigint_post_id = int(request_info.post_id)

                insert_comment_query = text(CommentUtil.insert_parent_comment())
                insert_comment_params = {
                    "comment_id": new_id,
                    "post_id": bigint_post_id,
                    "contents": request_info.contents,
                    "user_email": user_email,
                }
                s.execute(insert_comment_query, insert_comment_params)
                s.commit()

                # 본인이 작성한 거는 무시
                if request_info.post_author_email != user_email:
                    kafka_message = {
                        "url": f"{bigint_post_id}-{request_info.slug}?comment_id={new_id}",
                        "author_email": user_email,
                        "post_id": bigint_post_id,
                        "author_nickname": request_info.nickname,
                        "title": request_info.title,
                        "noti_type": "create_parent_comment",
                    }
                    json_str = json.dumps(kafka_message)
                    produce_notification(json_str)

                return {"result": 1}
        except Exception as e:
            logger.error(
                f"insert_parent_comment: {request_info.model_dump()}, error: {e}",
                exc_info=True,
            )
            return None

websocket.py

Websocket 구현 입니다.

Redis를 Listen 하는데 사용자 개별 연결을 email을 사용하여 진행합니다.

import asyncio
import json
import os
from typing import Dict, Set
from fastapi import WebSocket
from redis.asyncio import Redis
import logging

logger = logging.getLogger("ws")

# Redis 연결
redis = Redis.from_url(
    f"redis://{os.getenv('REDIS_HOST')}",
    decode_responses=True,
)

# In-memory (worker-local)
connected_websockets: Dict[str, WebSocket] = {}  # user_email -> websocket
user_listeners: Dict[str, asyncio.Task] = {}  # user_email -> redis listener task
sent_notifications: Dict[str, Set[str]] = {}  # user_email -> dedup key set


# WebSocket Handler
async def websocket_handler(websocket: WebSocket, user_email: str):
    """
    WebSocket 연결 처리
    """
    await websocket.accept()
    connected_websockets[user_email] = websocket
    logger.info(f"[OPEN] WebSocket 연결: {user_email}")

    # 1. 밀린 알림 먼저 전달 (Redis List)
    existing_notifications = []
    while True:
        n = await redis.lpop(f"notifications:{user_email}")
        if not n:
            break
        existing_notifications.append(json.loads(n))

    existing_notifications.reverse()

    await websocket.send_text(
        json.dumps(
            {
                "type": "init",
                "notifications": existing_notifications,
            }
        )
    )

    # 2. Redis listener 시작 (없으면)
    if user_email not in user_listeners:
        user_listeners[user_email] = asyncio.create_task(redis_listener(user_email))

    try:
        # ping / keep-alive 용
        while True:
            await websocket.receive_text()
    except Exception:
        pass
    finally:
        logger.info(f"[CLOSE] WebSocket 연결 종료: {user_email}")
        await cleanup_user(user_email)


# Redis Pub/Sub Listener
async def redis_listener(user_email: str):
    """
    Redis Pub/Sub 수신 → WebSocket 전송
    """
    pubsub = redis.pubsub()
    await pubsub.subscribe(f"notifications_channel:{user_email}")

    sent_notifications.setdefault(user_email, set())

    try:
        async for message in pubsub.listen():
            if message["type"] != "message":
                continue

            data = json.loads(message["data"])
            ws = connected_websockets.get(user_email)

            if not ws:
                # WS 없으면 그냥 소비 (유실 OK 정책)
                continue

            # WPF 위치 메시지
            if data.get("type") == "wpf_location":
                await ws.send_text(json.dumps(data))
                continue

            # 일반 알림 (중복 방지)
            notification_key = (
                f"{data.get('noti_type')}_"
                f"{data.get('post_id')}_"
                f"{data.get('parent_comment_id')}_"
                f"{data.get('author_email')}_"
                f"{data.get('id')}"
            )

            if notification_key in sent_notifications[user_email]:
                continue

            sent_notifications[user_email].add(notification_key)

            await ws.send_text(
                json.dumps(
                    {
                        "type": "message",
                        "data": data,
                    }
                )
            )

    except asyncio.CancelledError:
        pass
    finally:
        await pubsub.unsubscribe(f"notifications_channel:{user_email}")
        await pubsub.close()


# Cleanup
async def cleanup_user(user_email: str):
    """
    WebSocket / Listener / 메모리 정리
    """
    connected_websockets.pop(user_email, None)

    task = user_listeners.pop(user_email, None)
    if task:
        task.cancel()

    sent_notifications.pop(user_email, None)

main.py

구현한 handler를 main에 추가해줍니다.

from util.websocket import websocket_handler
from fastapi import FastAPI, WebSocket, Query

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = Query(...)):
    """
    로그인한 사용자만 WebSocket 연결
    토큰은 query param으로 전달
    """
    user_email = UserUtil.verify_google_token(token)
    if not user_email:
        await websocket.close(code=1008)  # 정책 위반
        return

    await websocket_handler(websocket, user_email)

NextJS 구현

NextJS 구현인데, 이 시기에는 사용자 알림만 사용해서 하나의 컴포넌트에만 연결을하면 되었기에 초기 코드를 보여드립니다.

use-weboscket.ts를 만들어 해당 컴포넌트에서 사용하는 방식입니다.

use-weboscket.ts

초기 연결시에는 알람 데이터를 그대로 사용하지만, 메시지가 추가로 올 경우에는 기존 데이터에 append 하는 방식 입니다.

import { useEffect } from "react";
import { wsStore } from "@/store/wsStore";
import { NotificationDataTypes } from "@/components/custom/NavBar/nav-bar.types";

export const useWebSocket = (accessToken?: string) => {
  const setNotifications = wsStore((s) => s.setNotifications);
  const setWpfLocation = wsStore((s) => s.setLocation);

  useEffect(() => {
    if (!accessToken) return;

    const ws = new WebSocket(
      `wss://${process.env.NEXT_PUBLIC_REDIS_HOST}/ws?token=${accessToken}`
    );

    ws.onmessage = (event) => {
      const parsed = JSON.parse(event.data);

      switch (parsed.type) {
        case "init": {
          const initial = parsed.notifications as NotificationDataTypes[];
          setNotifications(initial);
          break;
        }

        case "message": {
          const newNotification = parsed.data as NotificationDataTypes;
          setNotifications((prev = []) => [...prev, newNotification]);
          break;
        }
      }
    };

    return () => ws.close();
  }, [accessToken]);
};

nav-bar.tsx

상단에 고정으로 위치되어 있는 컴포넌트이고 해당 내용만 추가해줍니다.

import { useWebSocket } from "@/lib/hooks/useWebSocket";

useWebSocket(session?.accessToken);

확인

잘 되는지 확인해보면 됩니다.

저는 계정을 하나 더 만들어서 댓글을 달고 알림을 확인하며 검증했습니다.

겪은 문제점 및 대응

꽤 많은 우여곡절이 있었는데, 적어보겟습니다.

중복 메시지

중복 메시지가 많이 들어왔었습니다.

의도한 것은 하나씩만 이었는데, 실제로 해보니 엄청나게 많은 중복메시지가 들어왔습니다. 하나의 메시지가 생기면 3 ~ 4개 정도가 추가로 생겼습니다.

이것은 코드에서 파이썬의 set 데이터 구조를 사용하여 해결하였습니다.

sent_notifications: Dict[str, Set[str]] = {}  # user_email -> dedup key set

sent_notifications.setdefault(user_email, set())

코드 수정시 어려움

FastAPI 코드를 하나 수정하면 Websocket 전달이 잘 안되어서 문제를 겪었습니다.

에러 처리나 성능을 좀 개선하고 싶었는데, 의도처럼 안되어서 먹통이 된적이 꽤 있었습니다.

또한 확인도 늦어서 처리에 시간이 꽤 걸렸었습니다.

현재는 초기 작성 버전으로 유지 중입니다.

사용자가 알림을 읽었는지 여부 확인 어려움

이 부분 어렵습니다. 아직도 잘 안되어 있긴 하지만 간단하게는 조치를 해두었습니다.

알림에 마우스를 올린 것으로는 알림을 읽었다고 하기에는 어려울 것 같아 알림을 클릭하거나 마이 페이지에서 알림을 확인한 경우 해당 알림을 읽은 것으로 판단하고 있습니다.

그래서 클릭하거나, 마이 페이지의 알람으로 이동하면 DB에서 읽었는지 여부를 수정하고 있습니다.

후일담

Websocket을 처음 사용해봤는데, 생각처럼 안되는 부분도 있었고 매끄럽게 처리 되지 않는 부분도 있었습니다.

현재는 잘 사용하고 있습니다. 추후 작성할 외부 프로그램(WPF)와의 연동도 봐 주시기 바랍니다.