Swift Concurrency와 UICollectionView Prefetch

Layover 프로젝트를 진행하던 도중 재생 화면의 프로필 이미지와 위치 정보가 뒤늦게 나타나는 경우가 있어 이를 개선하고자 Prefetch를 적용시켜보았다.

UIKit에서 CollectionView나 TableView는 DataSourcePrefetching을 이용해 Prefetching을 지원한다.

https://developer.apple.com/documentation/uikit/uicollectionviewdatasourceprefetching/prefetching_collection_view_data

Prefetching collection view data | Apple Developer Documentation

collectionView.dataSource = dataSource

collectionView.prefetchDataSource = dataSource

// 이와 같은 방식도 가능
collectionView.prefetchDataSource = self

// UICollectionViewDataSourcePrefetching을 채택시켜 해당 메소드를 사용 가능하다.
// Prefetching이 필요한 시점에서 동작
extension ViewController: UICollectionViewDataSourcePrefetching {
    func collectionView(_ collectionView: UICollectionView, prefetchItemsAt indexPaths: [IndexPath]) {
				for indexPath in indexPaths {
            let model = models[indexPath.row]
            asyncFetcher.fetchAsync(model.identifier)
        }
    }

    func collectionView(_ collectionView: UICollectionView, cancelPrefetchingForItemsAt indexPaths: [IndexPath]) {
				for indexPath in indexPaths {
            let model = models[indexPath.row]
            asyncFetcher.cancelFetch(model.identifier)
        }
    }
}

위의 Prefetching CollectionView에서는 UICollectionViewDataSourcePrefetching을 채택시켜 prefetchItem의 indexPaths를 받아 prefetch한다.

기본적인 방식은 다음과 같다.

Main Thread에서 UICollectionViewDataSourcePrefetching와 같은 방식을 통해 prefetch가 Trigger되면 비동기적인 방식으로 Prefetch Thread에서 prefetch가 동작한다.

Main Thread에서 prefetch를 요청하고 Main Thread가 아닌 다른 Thread에서 prefetch를 진행해 해당 data를 Main Thread에서 쓴다.

필자는 여기서 궁금증이 하나 생겼는데, 그러면 Cache에는 아직 저장되지 않고, prefetching중인데 해당 데이터를 필요로 한다면?

Prefetch를 통해 데이터를 로드해 Cache에 저장하는데, 아직 로드 중일 때, Cache에 저장되지 않은 데이터를 사용해야할 순간이 오면 어떻게 될까? CollectionView의 경우 scroll을 내리다 보면 충분히 발생가능한 일이다.

Prefetch가 되지 않았으니 데이터 로드를 한번 더 한다. ←이는 너무 비효율적인 동작이다.

예제에서는 OperationQueue와 Completion Handler를 통해 이를 관리해주고 있었다.

하지만 Layover에서는 Swift Concurrency를 사용하고 있었기 때문에 이에 맞춰 변경을 해줄 필요가 있었다.

전체적인 흐름도를 보면 다음과 같다.

Task Dictionary안에 Prefetch Task들을 넣어 놓고 관리한다.

prefetching중인 데이터를 필요로 할 경우 id를 통해 Task Dictionary에 접근해서 해당 데이터 로드를 이어서 기다린다.

주요 코드는 다음과 같다.

private let imageCache: NSCache<NSURL, NSData> = {
        let cache = NSCache<NSURL, NSData>()
        // 10MB
        cache.totalCostLimit = 10 * 1024 * 1024
        return cache
    }()

private var fetchingImageTasks: [URL: Task<Data, Error>] = [:]
private let provider: ProviderType = Provider()

// PrefetchImage method URL을 key값으로 받아 사용한다.
func prefetchImage(for key: URL) async -> Data? {
	// imageCache에 해당 key가 있으면 cache에서 먼저 데이터를 가져온다.
    if let cachedData = imageCache.object(forKey: key as NSURL) as? Data {
        return cachedData
    }
	// cache에 데이터가 없으면 prefetching중인지 확인하고 있다면 해당 데이터를 기다린다.
    if let prefetchingData = fetchingImageTasks[key] {
        do {
            return try await prefetchingData.value
        } catch {
	          os_log(.error, log: .data, "%@", error.localizedDescription)
            return nil
        }
    }
      
	// 없을 경우 Task Dictionary에 key값과 Task동작을 넣어준다.
    let fetchTask: Task<Data, Error> = download(key)
    fetchingImageTasks[key] = fetchTask

	// data load 시작!
    do {
        let fetchedData: Data = try await fetchTask.value
        imageCache.setObject(fetchedData as NSData, forKey: key as NSURL)
        fetchingImageTasks[key] = nil
        return fetchedData
    } catch {
        os_log(.error, log: .data, "%@", error.localizedDescription)
        return nil
    }
}

// provider는 Layover의 Network를 담당.
// URLSession의 data method를 사용.
private func download(_ url: URL) -> Task<Data, Error> {
    return Task {
        try await provider.request(url: url)
	// 만약 사용한다면 이런식으로
	// let (data, _) = try await URLSession.shared.data(from: url)
        // return data
    }
}

전체 코드는 다음과 같다.

//
//  Prefetcher.swift
//  Layover
//
//  Created by 황지웅 on 1/18/24.
//  Copyright © 2024 CodeBomber. All rights reserved.
//

import Foundation
import CoreLocation
import OSLog

protocol PrefetchProtocol {
    func prefetchImage(for key: URL) async -> Data?
    func prefetchLocation(latitude: Double, longitude: Double) async -> String?
    func cancelPrefetchImage(for key: URL) async
    func cancelPrefetchLocation(latitude: Double, longitude: Double) async
}

actor Prefetcher: PrefetchProtocol {
    private let imageCache: NSCache<NSURL, NSData> = {
        let cache = NSCache<NSURL, NSData>()
        // 10MB
        cache.totalCostLimit = 10 * 1024 * 1024
        return cache
    }()

    private let locationCache: NSCache<NSString, NSString> = {
        let cache = NSCache<NSString, NSString>()
        // 5MB
        cache.totalCostLimit = 5 * 1024 * 1024
        return cache
    }()

    private var fetchingImageTasks: [URL: Task<Data, Error>] = [:]
    private var fetchingLocationTasks: [String: Task<String?, Error>] = [:]

    private let provider: ProviderType = Provider()

    func prefetchImage(for key: URL) async -> Data? {
        if let cachedData = imageCache.object(forKey: key as NSURL) as? Data {
            return cachedData
        }

        if let prefetchingData = fetchingImageTasks[key] {
            do {
                return try await prefetchingData.value
            } catch {
                os_log(.error, log: .data, "%@", error.localizedDescription)
                return nil
            }
        }
        
        let fetchTask: Task<Data, Error> = download(key)
        fetchingImageTasks[key] = fetchTask

        do {
            let fetchedData: Data = try await fetchTask.value
            imageCache.setObject(fetchedData as NSData, forKey: key as NSURL)
            fetchingImageTasks[key] = nil
            return fetchedData
        } catch {
            os_log(.error, log: .data, "%@", error.localizedDescription)
            return nil
        }
    }

    func prefetchLocation(latitude: Double, longitude: Double) async -> String? {
        let cacheKey = "\(latitude)-\(longitude)"

        if let cachedData = locationCache.object(forKey: cacheKey as NSString) as? String {
            return cachedData
        }

        if let prefetchingData = fetchingLocationTasks[cacheKey] {
            do {
                return try await prefetchingData.value
            } catch {
                os_log(.error, log: .data, "%@", error.localizedDescription)
                return nil
            }
        }

        let fetchTask: Task<String?, Error> = loadLocationInfo(latitude: latitude, longitude: longitude)
        fetchingLocationTasks[cacheKey] = fetchTask

        do {
            guard let fetchedData: String = try await fetchTask.value else { return nil }
            locationCache.setObject(fetchedData as NSString, forKey: cacheKey as NSString)
            return fetchedData
        } catch {
            os_log(.error, log: .data, "%@", error.localizedDescription)
            return nil
        }
    }

    func cancelPrefetchImage(for key: URL) async {
        fetchingImageTasks[key]?.cancel()
        fetchingImageTasks[key] = nil
    }

    func cancelPrefetchLocation(latitude: Double, longitude: Double) async {
        let key = "\(latitude)-\(longitude)"
        fetchingLocationTasks[key]?.cancel()
        fetchingLocationTasks[key] = nil
    }

    private func download(_ url: URL) -> Task<Data, Error> {
        return Task {
            try await provider.request(url: url)
        }
    }

    private func loadLocationInfo(latitude: Double, longitude: Double) -> Task<String?, Error> {
         return Task {
            let findLocation: CLLocation = CLLocation(latitude: latitude, longitude: longitude)
            let geoCoder: CLGeocoder = CLGeocoder()
            let localeIdentifier = Locale.preferredLanguages.first != nil ? Locale.preferredLanguages[0] : Locale.current.identifier
            let locale = Locale(identifier: localeIdentifier)
            let place = try await geoCoder.reverseGeocodeLocation(findLocation, preferredLocale: locale)
            return place.last?.administrativeArea
        }
    }
}

Layover에서는 재생 화면에서 프로필 이미지와 위치를 로드하기 때문에 두 경우를 Prefetch해서 사용했다.

여기서 중요한 점은 Actor를 사용한 것인데 class로 사용할 수도 있었겠지만 그럴 경우 data race가 발생할 수 있는 위험이 있다.

prefetchImage나 prefetchLocation의 경우 각 tasks dictionary에 접근을 하고 있는데 스크롤 속도가 빠른 숏폼 플랫폼의 특성상 동시에 접근할 수 있는 일이 생길 수 있다.

그래서 Actor를 사용하여 data race를 방지할 수 있는 코드로 작성하는 것이 중요하다.