Go로 PostgreSQL 프록시 만들기 (17) - Channel Blocking과 Connection Poisoning 버그 수정

Log() → eventCh(1024 버퍼) → processEvents 워커(1개) → handleEvent → sendWebhook

// 수정 전: handleEvent 내부
if l.httpClient != nil {
    l.sendWebhook(e)  // HTTP 호출 완료까지 워커가 멈춤
}

func (l *Logger) Log(e Event) {
    select {
    case l.eventCh <- e:
    default:
        slog.Warn("audit event channel full, dropping event")  // 여기로 빠짐
    }
}

// 수정 후
if l.httpClient != nil {
    go l.sendWebhook(e)  // 별도 고루틴으로 완전 분리
}

func (l *Logger) sendWebhook(e Event) {
    queryKey := truncateQuery(e.Query, 100)
    l.lastWebhookMu.Lock()
    if last, ok := l.lastWebhook[queryKey]; ok && time.Since(last) < l.webhookInterval {
        l.lastWebhookMu.Unlock()
        return  // 같은 패턴은 1분에 1번만
    }
    // ...
}

func TestAuditLogger_WebhookDoesNotBlockLogging(t *testing.T) {
    // 2초간 블로킹하는 mock webhook 서버
    ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        time.Sleep(2 * time.Second)
        w.WriteHeader(http.StatusOK)
    }))

    logger := New(cfg)

    // Slow Query → webhook 발동
    logger.Log(Event{DurationMS: 100, Query: "SELECT * FROM delay1"})
    time.Sleep(100 * time.Millisecond)

    // webhook이 블로킹 중인 동안 1100개 이벤트 전송
    for i := 0; i < 1100; i++ {
        logger.Log(Event{DurationMS: 5, Query: "SELECT * FROM fast"})
    }

    time.Sleep(500 * time.Millisecond)

    // 수정 전: 채널이 가득 차서 대부분 드랍
    // 수정 후: 프로세서가 블로킹되지 않아 채널이 즉시 소비됨
    remaining := len(logger.eventCh)
    if remaining > 10 {
        t.Errorf("processor was likely blocked, %d events remain", remaining)
    }
}

func (s *Server) executeOnPool(ctx context.Context, sql string, p *pool.Pool) (*QueryResponse, error) {
    conn, _ := p.Acquire(ctx)
    resp, execErr := executeQuery(conn, sql)
    if execErr != nil {
        p.Discard(conn)  // 쿼리 실패 시 정상적으로 폐기
        return nil, execErr
    }

    // 세션 리셋
    protocol.WriteMessage(conn, protocol.MsgQuery, resetPayload)
    drainUntilReady(conn)  // 여기서 에러가 나도 무시!
    p.Release(conn)        // 죽은 커넥션이 풀로 반환됨
    return resp, nil
}

func drainUntilReady(conn net.Conn) {
    for {
        msg, err := protocol.ReadMessage(conn)
        if err != nil || msg.Type == protocol.MsgReadyForQuery {
            return  // 에러든 성공이든 그냥 리턴
        }
    }
}

// drainUntilReady가 성공/실패를 구분하도록 변경
func drainUntilReady(conn net.Conn) error {
    for {
        msg, err := protocol.ReadMessage(conn)
        if err != nil {
            return err  // 에러 시 실패 반환
        }
        if msg.Type == protocol.MsgReadyForQuery {
            return nil  // 성공적으로 ReadyForQuery 수신
        }
    }
}

// 수정 후
if err := drainUntilReady(conn); err != nil {
    p.Discard(conn)   // 죽은 커넥션은 영구 폐기
    return resp, nil  // 쿼리 결과는 이미 받았으므로 정상 반환
}
p.Release(conn)  // 정상 커넥션만 풀에 반환

func TestDrainUntilReady_ReturnsErrorOnBrokenConn(t *testing.T) {
    conn := &brokenConn{}  // Read()가 항상 net.ErrClosed 반환
    err := drainUntilReady(conn)
    if err == nil {
        t.Fatal("expected error from drainUntilReady on broken connection, got nil")
    }
}