[Loki] Shard Factor 1일 때 Shuffle Shard 생략으로 메모리 50% 절감

PR 링크: grafana/loki#21190 상태: Merged | 변경: +27 / -104

들어가며

Grafana Loki의 distributor에서 segmentation key 기반으로 Kafka 파티션을 결정할 때, ShuffleShard() 호출로 서브링을 생성합니다. 그런데 shard factor가 1인 경우(단일 파티션만 필요)에도 ShuffleShard()를 호출하여 불필요한 연산과 메모리 캐싱이 발생했습니다. 이 PR은 shard factor가 1이면 ActivePartitionForKey()로 직접 파티션을 결정하도록 최적화했으며, 실제 운영 환경에서 메모리 사용량이 50% 감소했습니다.

핵심 코드 분석

Before: 항상 ShuffleShard 호출

func (r *segmentationPartitionResolver) Resolve(...) (int32, error) {
    // rate가 0이면 랜덤 파티션 선택
    if rateBytes == 0 {
        activePartitionIDs := subring.ActivePartitionIDs()
        rand.Shuffle(len(activePartitionIDs), func(i, j int) {
            activePartitionIDs[i], activePartitionIDs[j] = activePartitionIDs[j], activePartitionIDs[i]
        })
        return activePartitionIDs[0], nil
    }
    subring, err = r.getSegmentationKeySubring(ctx, subring, key, rateBytes)
    // ...
    return subring.ActivePartitionForKey(hashKey)
}

After: shard factor 1이면 ShuffleShard 건너뛰기

func (r *segmentationPartitionResolver) Resolve(...) (int32, error) {
    // rate가 0이면 키 기반으로 파티션 결정
    if rateBytes == 0 {
        return subring.ActivePartitionForKey(uint32(key.Sum64()))
    }
    numShuffleShardPartitions := numPartitionsForRate(rateBytes, r.perPartitionRateBytes, subring.ActivePartitionsCount())
    // shard가 1개면 ShuffleShard 불필요
    if numShuffleShardPartitions == 1 {
        return subring.ActivePartitionForKey(uint32(key.Sum64()))
    }
    subring, err = subring.ShuffleShard(string(key), numShuffleShardPartitions)
    // ...
    return subring.ActivePartitionForKey(hashKey)
}

추가로, rate가 0인 경우의 랜덤 셔플도 제거하고 ActivePartitionForKey()를 사용하여 결정적(deterministic) 파티션 할당으로 변경했습니다.

왜 이게 좋은가

1. 메모리 50% 절감: 실제 운영에서 segmentation key별 shuffle shard 캐시가 상당한 메모리를 사용했는데, shard factor 1인 경우(대부분의 작은 segmentation key)의 캐시가 완전히 제거됩니다.
2. CPU 절감: ShuffleShard() 내부의 해싱, 정렬, 서브링 생성 연산을 건너뜁니다.
3. 코드 단순화: getSegmentationKeySubring 함수가 제거되고, randomlySharded 메트릭도 불필요해져 삭제됩니다. 104줄이 삭제되고 27줄만 추가된 이유입니다.
4. 결정적 동작: 랜덤 셔플 대신 해시 기반 파티션 결정으로 같은 키가 항상 같은 파티션에 매핑되어 데이터 지역성이 향상됩니다.

조건부 최적화의 좋은 사례입니다. "shard 1개에 ShuffleShard가 필요한가?"라는 단순한 질문이 실제 운영 환경에서 메모리 50% 절감이라는 큰 효과를 가져왔습니다.

참고 자료

• Grafana Loki Shuffle Sharding
• Consistent Hashing과 파티션 할당