В данной части из серии swift concurrency мы продолжим исследовать structured concurrency. В этот раз остановимся на сущности под названием TaskGroup. Узнаем, как с ней работать и чем она отличается от Task и async let. На примере сравним аналогичные реализации с и без TaskGroup и разберем некоторые (не для всех очевидные) моменты при работе с данной сущностью.

Статьи из серии

Оглавление

Structured concurrency
Сетапим окружение для примера
Реализуем загрузку N изображений (без TaskGroup)
Загрузка N изображений через TaskGroup
Нюансы при работе с TaskGroup
Итоги
Полезные ссылки

Structured concurrency

В прошлой статье мы начали знакомиться с структурным параллелизмом (structured concurrency) и частью его инструментов. Эта статья является логическим ее продолжением, но в любом случае кратко напомню суть:

Structured concurrency - это раздел из Swift concurrency, который задает и реализует ряд механизмов и объектов, позволяющих конкуррентно (одновременно/параллельно) выполнять какую-либо работу, выстраивать зависимости между выполняемыми задачами и гибко управлять их иерархией (например, структурно отменять определенную группу задач).

Сетапим окружение для примера

Материал лучше всего осваивается на примерах. Предлагаю не медлить и начать сразу с него. Продолжим докручивать функциональностью пример из предыдущей статьи. Там мы реализовали контроллер, который умеет загружать картинку при нажатии на кнопку (rocket science).

Напомню, что для запросов в сеть мы пользуемся публичным сервисом jsonplaceholder. Код из предыдущей статьи:

Photo

import Foundation

// Структура для хранения ответа сервера
struct Photo: Decodable {
  let albumId: Int
  let id: Int
  let title: String
  let url: URL
  let thumbnailUrl: URL
}

ImageLoader

import UIKit

// Класс, отвечающий за загрузку изображений. В прошлых статьях
// использовалось computed property для получения изображения.
// Тут переписал на обычный метод для удобства

enum ImageLoaderError: Error {
  case incorrectImageData
}

final class ImageLoader: Sendable {
  func loadImage(from url: URL) async throws -> UIImage {
    let (data, _) = try await URLSession.shared.data(from: url)
    guard let image = UIImage(data: data) else {
      throw ImageLoaderError.incorrectImageData
    }
    return image
  }
}

ViewController

import UIKit

enum ViewControllerError: Error {
  case emptyData
}

// Контроллер, который отображает кнопку и картинку. 
// При нажатии на кнопку улетает запрос к jsonplaceholder 
// на получение метаданных изображений. Из метаданных берется 
// первый url и загружается с помощью ImageLoader

final class ViewController: UIViewController {
  private let imageView: UIImageView = {
    let view = UIImageView()
    view.translatesAutoresizingMaskIntoConstraints = false
    return view
  }()

  private let button: UIButton = {
    let button = UIButton(type: .system)
    button.setTitle("Load image", for: .normal)
    button.translatesAutoresizingMaskIntoConstraints = false
    return button
  }()

  private let imageLoader = ImageLoader()

  private var imageShowingTask: Task<Void, Never>?

  init() {
    super.init(nibName: nil, bundle: nil)
  }

  required init?(coder: NSCoder) {
    fatalError("init(coder:) has not been implemented")
  }

  override func viewDidLoad() {
    super.viewDidLoad()
    setupView()
  }

  override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    imageShowingTask?.cancel()
  }

  private func onLoadImageButtonTap() {
    guard imageShowingTask == nil else { return }

    imageShowingTask = Task {
      do {
        let image = try await loadImage()
        imageView.image = image
      } catch {
        imageShowingTask = nil
      }
    }
  }

  private func loadImage() async throws -> UIImage {
    let photos = try await getPhotos()
    guard !photos.isEmpty else { throw ViewControllerError.emptyData }

    return try await imageLoader.loadImage(from: photos[0].url)
  }

  func getPhotos() async throws -> [Photo] {
    let url = URL(string: "https://jsonplaceholder.typicode.com/albums/1/photos")!
    let request = URLRequest(url: url)

    let (data, _) = try await URLSession.shared.data(for: request)

    let photos = try JSONDecoder().decode([Photo].self, from: data)
    return photos
  }

  private func setupView() {
    view.backgroundColor = .white
    button.addAction(
      UIAction { [weak self] _ in
        self?.onLoadImageButtonTap()
      },
      for: .touchUpInside
    )

    addSubviews()
    setupConstraints()
  }

  private func addSubviews() {
    view.addSubview(imageView)
    view.addSubview(button)
  }

  private func setupConstraints() {
    NSLayoutConstraint.activate([
      imageView.centerYAnchor.constraint(equalTo: view.centerYAnchor),
      imageView.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      imageView.widthAnchor.constraint(equalToConstant: 248),
      imageView.heightAnchor.constraint(equalToConstant: 248),

      button.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      button.topAnchor.constraint(equalTo: imageView.bottomAnchor, constant: 24),
    ])
  }
}

Менеджеры выяснили, что пользователям скучно нажимать на кнопку и видеть после этого одну и ту же картинку, поэтому они пришли к нам с запросом отображать рандомную картинку при нажатии. Помимо этого, они выяснили, что подавляющее большинство юзеров будет тыкать на кнопку много раз, и хочется, чтобы наше приложение умело обновлять изображения как можно быстрее. Давайте реализуем эту киллер-фичу.

Исходя из требований можно понять, что для реализации минимальной задержки нам нужно предзагрузить пул изображений, и при нажатии на кнопку выбирать случайное из этого пулла и сразу же его отображать (не теряя время каждый раз на загрузку отдельного изображения).

Для начала немного преобразуем ImageLoader. Нам требуется, чтобы он умел грузить как одно, так и несколько изображений за раз.

enum ImageLoaderError: Error {
  case incorrectImageData
  // Также добавим новый тип ошибки, который понадобится в дальнейшем
  case emptyResult
}

final class ImageLoader {
  func loadImage(from url: URL) async throws -> UIImage {
    // На момент написания статьи API выдает 504 Gateway Time Out на некоторые
    // изображения, поэтому немного обновим функцию, добавив таймаут
    let request = URLRequest(url: url, timeoutInterval: 1)
    let (data, _) = try await URLSession.shared.data(for: request)
    guard let image = UIImage(data: data) else {
      throw ImageLoaderError.incorrectImageData
    }
    return image
  }

  // Новый метод по загрузке массива изображений
  func loadImages(from urls: [URL]) async throws -> [UIImage] {
    ...
  }
}

Ориентируясь на сигнатуру функции, давайте сразу прикрутим необходимую логику к контроллеру, чтобы у вас была возможность в процессе написания нового метода loadImages запускать код и смотреть, как он себя ведет:

Обновленный ViewController

import UIKit

enum ViewControllerError: Error {
  case emptyData
}

final class ViewController: UIViewController {
  private let imageView: UIImageView = {
    let view = UIImageView()
    view.translatesAutoresizingMaskIntoConstraints = false
    return view
  }()

  private let button: UIButton = {
    let button = UIButton(type: .system)
    button.setTitle("Load image", for: .normal)
    button.translatesAutoresizingMaskIntoConstraints = false
    return button
  }()

  private let imageLoader = ImageLoader()
  
  // 1
  private var imagesLoadingTask: Task<[UIImage], Never>?
  private var imageShowingTask: Task<Void, Never>?

  init() {
    super.init(nibName: nil, bundle: nil)
  }

  required init?(coder: NSCoder) {
    fatalError("init(coder:) has not been implemented")
  }

  override func viewDidLoad() {
    super.viewDidLoad()
    setupView()
  }

  override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    // 7
    imagesLoadingTask?.cancel()
    imageShowingTask?.cancel()
  }

  private func onLoadImageButtonTap() {
    guard imageShowingTask == nil else { return }

    imageShowingTask = Task {
      // 4
      let loadingTask = imagesLoadingTask ?? startNewLoadingTask()
      let images = await loadingTask.value
      imageView.image = images.randomElement()
      // 5
      imageShowingTask = nil
    }
  }

  // 2
  private func startNewLoadingTask() -> Task<[UIImage], Never> {
    let newTask = Task {
      do {
        return try await loadImages()
      } catch {
        // 3
        imagesLoadingTask = nil
        return []
      }
    }
    imagesLoadingTask = newTask
    return newTask
  }
  
  // 6
  private func loadImages() async throws -> [UIImage] {
    let photos = try await getPhotos()
    return try await imageLoader.loadImages(from: photos.map(\.url))
  }

  func getPhotos() async throws -> [Photo] {
    let url = URL(string: "https://jsonplaceholder.typicode.com/albums/1/photos")!
    let request = URLRequest(url: url)

    let (data, _) = try await URLSession.shared.data(for: request)

    let photos = try JSONDecoder().decode([Photo].self, from: data)
    return photos
  }

  private func setupView() {
    view.backgroundColor = .white
    button.addAction(
      UIAction { [weak self] _ in
        self?.onLoadImageButtonTap()
      },
      for: .touchUpInside
    )

    addSubviews()
    setupConstraints()
  }

  private func addSubviews() {
    view.addSubview(imageView)
    view.addSubview(button)
  }

  private func setupConstraints() {
    NSLayoutConstraint.activate([
      imageView.centerYAnchor.constraint(equalTo: view.centerYAnchor),
      imageView.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      imageView.widthAnchor.constraint(equalToConstant: 248),
      imageView.heightAnchor.constraint(equalToConstant: 248),

      button.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      button.topAnchor.constraint(equalTo: imageView.bottomAnchor, constant: 24),
    ])
  }
}

Теперь храним две Task: imagesLoadingTask для загрузки и imageShowingTask для отображения. В рамках imagesLoadingTask осуществляется загрузка массива изображений и дальнейшее получение этого массива без повторной загрузки. imageShowingTask - это старая Task, в рамках которой осуществляется присвоение изображения в UIImageView. Держим ее в классе, чтобы быть уверенными в том, что в момент времени выполняется только одна Task. Делается это для того, чтобы пользователь не наплодил кучу тасок быстрыми нажатиями на кнопку загрузки.
Новый метод, который создает и возвращает загрузочную таску.
Не забываем про обработку ошибок. Если метод loadImages вернет ошибку, то мы зануляем таску, что позволит при последующем нажатии на кнопку создать новую загрузочную таску.
Проверяем, существует ли загрузочная таска. Если да, то используем ее. Она может быть в 2-х состояниях: либо все еще выполняется, тогда с помощью await value мы дождемся результата, либо уже завершена, тогда мы просто достанем из нее результат. Если таски не существует, то инициируем ее создание с помощью метода startNewLoadingTask().
Зануляем таску показа изображения после завершения. Это позволит снова ее создать при последующем нажатии на кнопку и отобразить новое изображение.
Метод теперь грузит не одно изображение, а сразу массив.
Не забываем отменять таски при уходе с экрана. Иначе они будут жить, пока не закончат выполнение, и держать контроллер в памяти.

Реализуем загрузку N изображений (без TaskGroup)

Окружение для нашего загрузчика изображений настроено, теперь перейдем к реализации loadImages. Начнем реализовывать уже известными для нас способами:

func loadImages(from urls: [URL]) async throws -> [UIImage] {  
  var images: [UIImage] = []
  for url in urls {
    // 1
    guard let image = try? await loadImage(from: url) else { continue }
    images.append(image)
  }
  // 2
  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  return images
}

Тут важный нюанс. Так как мы загружаем N изображений, нецелесообразно завершать всю таску, если загрузка одного изображения выбросит ошибку, поэтому будем игнорировать проброс ошибки с помощью try? и пропускать изображение, если по какой-то причине загрузка вернет ошибку.
Для нас не критична ошибка при загрузке нескольких изображений. Но если не удалось загрузить все - будем считать ошибкой такой сценарий.

Метод отработает как надо, только в нем есть проблема: на каждой итерации цикла мы ждем, пока одно изображение загрузится. И так для каждого URL из массива. То есть в любой момент времени у нас грузится только одно изображение, хотя было бы быстрее загружать их конкуррентно.

Каждая Task может в моменте выполнять только одну функцию. Но что нам мешает реализовать отдельную Task для загрузки каждого изображения из массива?

func loadImages(from urls: [URL]) async throws -> [UIImage] {
  // 1
  var loadingTasks: [Task<UIImage, Error>] = []
  for url in urls {
    loadingTasks.append(Task {
      try await loadImage(from: url)
    })
  }

  // 2
  var images: [UIImage] = []
  for task in loadingTasks {
    guard let image = try? await task.value else { continue }
    images.append(image)
  }

  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  
  return images
}

Создаем массив тасок и заполняем его, создавая отдельную таску для загрузки каждого URL. По мере заполнения массива таски начинают загрузку (конкуррентно).
После заполнения ожидаем результат каждой из тасок в цикле.

Если вы запустите прошлый вариант (синхронный), а потом этот, то сразу увидите прирост в скорости.

Но и у этого варианта есть проблема. В случае, если нам нужно будет отменить загрузку, ничего не произойдет, и функция выполнится как если бы мы не вызывали функцию cancel. Проверить этот факт можно, немного видоизменив метод ViewController.loadImages следующим образом:

private func loadImages() async throws -> [UIImage] {
  let photos = try await getPhotos()
  // 1
  let imagesLoadingTask = Task {
    try await imageLoader.loadImages(from: photos.map(\.url))
  }
  // 2
  imagesLoadingTask.cancel()
  return try await imagesLoadingTask.value
}

Оборачиваем загрузку в другую Task, чтобы можно было проверять функциональность отмены только у ImageLoader, так как если бы мы отменяли таску целиком, то ошибка бы выбросилась, не доходя до нашей функции.
Отменяем эту Task сразу же после создания.

Теперь расставим логи в нашей функции загрузки:

func loadImages(from urls: [URL]) async throws -> [UIImage] {  
  var loadingTasks: [Task<UIImage, Error>] = []
  for url in urls {
    loadingTasks.append(Task {
      
      print("[Inside Task] Start loading image from url \(url.absoluteString)")
      print("[Inside Task] Task.isCancelled = \(Task.isCancelled)")
      
      let image = try await loadImage(from: url)
      
      print("[Inside Task] Success loading image from url \(url.absoluteString)")
      print("[Inside Task] Task.isCancelled = \(Task.isCancelled)")
      
      return image
    })
  }
  
  var images: [UIImage] = []

  print("[loadImages] Task.isCancelled before wait loop: \(Task.isCancelled)")
  for task in loadingTasks {
    guard let image = try? await task.value else { continue }
    images.append(image)
  }

  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  print("[loadImages] Success load \(images.count) images. Task.isCancelled = \(Task.isCancelled)")
  return images
}

Если запустить этот код, то в логах можно будет увидеть что-то около:

[Inside Task] Start loading image from url https://via.placeholder.com/600/92c952
[Inside Task] Task.isCancelled = false
[Inside Task] Start loading image from url https://via.placeholder.com/600/771796
[Inside Task] Task.isCancelled = false
...
[loadImages] Task.isCancelled before wait loop: true
...
[Inside Task] Success loading image from url https://via.placeholder.com/600/771796
[Inside Task] Task.isCancelled = false
...
[loadImages] Success load 6 images. Task.isCancelled = true

Task.isCancelled выставлено в true в области функции loadImages, и несмотря на это цикл продолжает крутиться и изображения продолжают грузиться. Внутри самих Task isCancelled будет равен false. Это происходит по двум причинам:

Task, которые мы создаем внутри функции, не являются child по отношению к текущей, и из-за этого при отмене они не помечаются автоматически флагом cancelled.
Мы не заложили логику отмены в текущую функцию. Напомню, что вызов метода Task.cancel() всего лишь помечает Task как canceled (буквально bool переменную в true выставляет, ничего более). Разработчикам требуется следить за этим флагом внутри функции и реализовывать логику отмены, если он в какой-то момент выставится в true.

Как уже упоминалось выше, любая асинхронная функция выполняется в рамках какой-то таски, и любая таска внутри не конкуррентна. Функция loadImages выполняется в рамках условной Task1, при отмене мы пометим как cancelled только ее, но не все остальные:

Следовательно, механизм отмены всей иерархии Task's из structured concurrency в нашем случае не заработает "из коробки", потребуются дополнительные действия для поддержки этого функционала:

func loadImages(from urls: [URL]) async throws -> [UIImage] {  
  var loadingTasks: [Task<UIImage, Error>] = []

  // 1
  try Task.checkCancellation()
//  guard !Task.isCancelled else { throw CancellationError() }

  for url in urls {
    loadingTasks.append(Task {
      
      print("[Inside Task] Start loading image from url \(url.absoluteString)")
      print("[Inside Task] Task.isCancelled = \(Task.isCancelled)")
      
      let image = try await loadImage(from: url)
      
      print("[Inside Task] Success loading image from url \(url.absoluteString)")
      print("[Inside Task] Task.isCancelled = \(Task.isCancelled)")
      
      return image
    })
  }

  var images = [UIImage]()

  // 2
  do {
    print("[loadImages] Task.isCancelled before wait loop: \(Task.isCancelled)")
    for task in loadingTasks {
      // 3
      try Task.checkCancellation()
      // 4
      guard let image = try? await task.value else { continue }
      images.append(image)
    }
  } catch {
    // 5
    loadingTasks.forEach { $0.cancel() }
    throw error
  }

  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  print("[loadImages] Success loaded \(images.count) images. Task.isCancelled = \(Task.isCancelled)")
  return images
}

Перед созданием массива тасок проверяем, не отменена ли Task, в рамках которой мы выполняемся. Напомню, что проверить это можно двумя способами: через checkCancellation или через isCancelled. Первый вариант более удобный, так как он сразу выбрасывает CancellationError() в случае, если isCancelled == true. В закомментированном коде указал эквивалентную конструкцию при помощи isCancelled.
Оборачиваем ожидающий цикл в do/catch. Это делается для того, чтобы у нас была возможность отреагировать на ошибку из блока do в блоке catch. Без этого ошибка бы просто пробросилась выше, и у нас не было бы возможности руками отменить все созданные Task's. В предыдущем пункте мы не обернули проверку на отмену в do/catch из-за того, что массив тасок пуст и отменять на тот момент нечего.
Перед каждым await проверяем, не отменена ли Task, в рамках которой мы выполняемся.
Загрузку конкретного изображения намеренно оставляем с try?, так как мы не хотим завершать загрузку всех из-за ошибки в одной.
В блоке catch отменяем все созданные Task's и выбрасываем ошибку выше. В текущей реализации мы можем попасть сюда только из-за ошибки отмены, но это гибко расширяемое решение. И если мы, к примеру, захотим останавливать все таски при ошибке загрузки одного из изображений, то нам достаточно поменять try? на try, и ожидаемое поведение заработает как нужно.

После запуска функции в таком виде консоль будет пуста, что говорит о том, что никакая лишняя работа не была запущена. Мы отменяем задачу моментально после создания, но даже если отменять Task в самом разгаре загрузки изображений - она успешно остановит всю работу (как только наткнется на любую из наших проверок на isCancelled). Можете поэкспериментировать и добавить для отмены какой-либо delay и проверить логи после этого.

Возможно у некоторых возникнет вопрос: в функции loadImage, которая отвечает за загрузку одного изображения, мы тоже не проверяем флаг isCancelled, получается, отмена не работает внутри нее?

func loadImage(from url: URL) async throws -> UIImage {
  let request = URLRequest(url: url, timeoutInterval: 1)
  let (data, _) = try await URLSession.shared.data(for: request)
  guard let image = UIImage(data: data) else {
    throw ImageLoaderError.incorrectImageData
  }
  return image
}

В данном случае отмена работает как надо, и поддерживать ее руками не нужно. Все из-за того, что нативная функция URLSession.shared.data внутри себя реализует логику проверки и выброса ошибки в случае отмены.

Реализацию выше можно сделать (и сильно упростить) при помощи "коробочного" решения - TaskGroup.

Загрузка N изображений через TaskGroup

TaskGroup - это сущность, с помощью которой можно создать и запустить конкуррентно N child Task's. С помощью конструкции async let можно создать только заранее известное количество Task's. При создании unstructured Tasks (Task {}) можно решать о количестве задач в рантайме, но они не будут child по отношению к текущей. TaskGroup как раз позволяет объединить озвученные преимущества 2-х подходов по созданию Task. Они будут child, и их количество определяется в runtime.

Теперь давайте перепишем нашу функцию, используя этот инструмент:

func loadImages(from urls: [URL]) async throws -> [UIImage] {  
  // 1
  let images = await withTaskGroup(of: UIImage?.self, returning: Array<UIImage>.self) { group in
    var images = [UIImage]()

    for url in urls {
      // 2
      group.addTask {
        // 3
        try? await self.loadImage(from: url)
      }
    }

    // 4                                               
    for await image in group {
      guard let image else { continue }
      images.append(image)
    }                                                                                         

    return images                                                                                   
  }

  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  return images
}

TaskGroup инициализируются с помощью конструкции withTaskGroup (либо withThrowingTaskGroup для групп, которые могут выбросить ошибку). Данная функция принимает тип, возвращаемый каждой child Task из группы. В нашем случае это UIImage?, так как каждая таска по отдельности может вернуть экземпляр изображения или nil. Вторым параметром идет тип, который возвращает вся группа (его можно не указывать явно, компилятор поймет и без этого), ну и завершающим параметром идет замыкание-обработчик, в которое передается сама группа и в рамках которого необходимо реализовать логику по работе с этой группой (создание/обработка childTask's).
Child Task's создаются через метод addTask у группы. Еще раз отмечу, что это именно child task's, следовательно они попадают в иерархию тасок и автоматически помечаются как отмененные в случае, когда отменяется задача выше по иерархии.
Создавая TaskGroup через withTaskGroup, мы не можем выбрасывать ошибки (так как она не throws), поэтому на уровне child task's останавливаем проброс ошибки выше с помощью конструкции try?. Вместо ошибки у нас вернется nil, которую мы отфильтруем на следующем этапе.
TaskGroup имплементирует протокол AsyncSequence, поэтому можно работать с ней через for _ in, или через функции высших порядков (filter, compactMap, reduce).

Вариант с помощью функций высших порядков будет выглядеть так (вместо for await цикла и хранения массива UIImage):

return await group
  .compactMap { $0 }
  .reduce(into: Array<UIImage>()) { $0.append($1) }

Но в текущей реализации оставлю вариант с for await. Он немного проще воспринимается при чтении кода.

Пока у нас есть подготовленная среда для проверки отмены - грех это не сделать. Перед этим расставим логи:

func loadImages(from urls: [URL]) async throws -> [UIImage] {
  let images = await withTaskGroup(of: UIImage?.self, returning: Array<UIImage>.self) { group in
    var images = [UIImage]()

    for url in urls {
      group.addTask {
        print("[Subtask] Start loading image from url \(url.absoluteString)")
        print("[Subtask] Task.isCancelled = \(Task.isCancelled)")

        // Вместо try? обернем вызов в try/catch. Чтобы явно увидеть по логам,
        // что загрузка не завершится в случае отмены
        do {
          let image = try await self.loadImage(from: url)
          print("[Subtask] Success loading image from url \(url.absoluteString)")
          print("[Subtask] Task.isCancelled = \(Task.isCancelled)")
          return image
        } catch {
          return nil
        }
      }
    }

    for await image in group {
      guard let image else { continue }
      images.append(image)
    }

    print("[TaskGroup] Success loaded \(images.count) images. Task.isCancelled = \(Task.isCancelled)")
    return images
  }

  guard !images.isEmpty else { throw ImageLoaderError.emptyResult }
  return images
}

Запустив эту функцию, в консоль выведется примерно следующее:

[Subtask] Start loading image from url https://via.placeholder.com/600/771796
[Subtask] Task.isCancelled = true
...
[TaskGroup] Success loaded 0 images. Task.isCancelled = true

В последнем логе мы увидим, что было успешно загружено 0 изображений. Следовательно, все загрузочные таски отменились сразу же после создания. Отмена заложена в TaskGroup автоматически, с нашей стороны никаких дополнительных проверок и выбросов ошибок в этом варианте реализации не требуется.

Но перед последним логом мы видим N сообщений о том, что загрузка началась. Возникает резонный вопрос: "А зачем их вообще создавать, если уже в начале выполнения TaskGroup было известно, что она отменена?". И у Apple есть решение. В момент добавления тасок в группу можно воспользоваться вместо group.addTask методом group.addTaskUnlessCancelled. По названию можно догадаться, что метод не сработает в случае, если Task, в рамках которой выполняется группа, отменена. Но если вы просто поменяется вызов addTask на addTaskUnlessCancelled, то ваши глаза начнет резать warning. Все дело в том, что этот метод возвращает флаг, сигнализирующий о том, создалась ли по итогу childTask (и он не @discardableResult). Конкретно в нашем случае это излишняя информация, поэтому можно написать расширение с методом, который будет явно игнорировать возвращаемое значение:

fileprivate extension TaskGroup {
  @discardableResult
  mutating func addTaskUnlessCancelledDiscardableResult(
    priority: TaskPriority? = nil,
    operation: sending @escaping @isolated(any) () async -> ChildTaskResult
  ) -> Bool {
    addTaskUnlessCancelled(priority: priority, operation: operation)
  }
}

При замене addTask на addTaskUnlessCancelled (и запуске) вы увидите один единственный лог в консоли, который можно повесить в рамочку и гордится тем, что при отмене наша группа не выполняет никакой лишней работы.

[TaskGroup] Success loaded 0 images. Task.isCancelled = true

Для завершения нашего примера остался один нюанс. В текущей реализации изображения грузятся конкуррентно, но контроллер все равно ждет загрузки всех изображений (так как мы ждем завершения всей группы). Давайте усовершенствуем загрузчик при помощи AsyncStream:

// 1
func loadImages(from urls: [URL]) async throws -> AsyncThrowingStream<UIImage, Error> {
  let (stream, continuation) = AsyncThrowingStream<UIImage, Error>.makeStream()

  // 2
  let task = Task {
    await withTaskGroup(of: UIImage?.self) { group in
      for url in urls {
        group.addTaskUnlessCancelledDiscardableResult {
          try? await self.loadImage(from: url)
        }
      }

      // 3
      var hasImages = false
      for await image in group {
        guard let image else { continue }
        continuation.yield(image)
        hasImages = true
      }

      guard hasImages else {
        continuation.finish(throwing: ImageLoaderError.emptyResult)
        return
      }
                                        
      // 4
      continuation.finish()
    }
  }

  // 5
  continuation.onTermination = { _ in
    task.cancel()
  }

  return stream
}

Вместо массива UIImage будем возвращать stream, с помощью которого можно будет получать изображения по мере их загрузки.
Создаем unstructured Task, так как мы не хотим ждать завершения TaskGroup в рамках выполнения функции (будем ждать завершения в рамках этой Task).
Теперь нам не нужно хранить массив UIImage. По мере загрузки будем прокидывать изображения в stream. Флаг hasImages нужен, чтобы после цикла убедиться, что мы загрузили как минимум 1 изображение.
Не забываем завершать stream.
Так как TaskGroup выполняется в рамках unstructured Task, отмена не будет прокидываться до нее автоматически. Для отмены воспользуемся замыканием onTermination у continuation. Которое вызывается при отмене таски, в рамках которой выполняется ожидание стрима.

Контроллер тоже необходимо немного модифицировать под новую функцию (и убрать из него тестовый код с отменой, добавленный ранее):

Контроллер с AsyncStream

Обновленные кусочки помечены цифрами

import UIKit

enum ViewControllerError: Error {
  case emptyData
}

final class ViewController: UIViewController {
  private let imageView: UIImageView = {
    let view = UIImageView()
    view.translatesAutoresizingMaskIntoConstraints = false
    return view
  }()

  private let button: UIButton = {
    let button = UIButton(type: .system)
    button.setTitle("Load image", for: .normal)
    button.translatesAutoresizingMaskIntoConstraints = false
    return button
  }()

  private let imageLoader = ImageLoader()

  // 2
  private var images: [UIImage] = []
  private var imagesLoadingTask: Task<Void, Never>?

  init() {
    super.init(nibName: nil, bundle: nil)
  }

  required init?(coder: NSCoder) {
    fatalError("init(coder:) has not been implemented")
  }

  override func viewDidLoad() {
    super.viewDidLoad()
    setupView()
  }

  override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    imagesLoadingTask?.cancel()
  }

  // 5
  private func onLoadImageButtonTap() {
    startNewLoadingTaskIfNeeded()
    if let image = images.randomElement() {
      imageView.image = image
    }
  }

  private func startNewLoadingTaskIfNeeded() {
    guard imagesLoadingTask == nil else { return }

    imagesLoadingTask = Task {
      do {
        let stream = try await startLoadingImages()

        // 3
        for try await image in stream {
          images.append(image)
          if imageView.image == nil {
            imageView.image = image
          }
        }
      } catch {
        // 4
        imagesLoadingTask = nil
      }
    }
  }

  // 1
  private func startLoadingImages() async throws -> AsyncThrowingStream<UIImage, Error> {
    let photos = try await getPhotos()
    return try await imageLoader.loadImages(from: photos.map(\.url))
  }

  func getPhotos() async throws -> [Photo] {
    let url = URL(string: "https://jsonplaceholder.typicode.com/albums/1/photos")!
    let request = URLRequest(url: url)

    let (data, _) = try await URLSession.shared.data(for: request)

    let photos = try JSONDecoder().decode([Photo].self, from: data)
    return photos
  }

  private func setupView() {
    view.backgroundColor = .white
    button.addAction(
      UIAction { [weak self] _ in
        self?.onLoadImageButtonTap()
      },
      for: .touchUpInside
    )

    addSubviews()
    setupConstraints()
  }

  private func addSubviews() {
    view.addSubview(imageView)
    view.addSubview(button)
  }

  private func setupConstraints() {
    NSLayoutConstraint.activate([
      imageView.centerYAnchor.constraint(equalTo: view.centerYAnchor),
      imageView.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      imageView.widthAnchor.constraint(equalToConstant: 248),
      imageView.heightAnchor.constraint(equalToConstant: 248),

      button.centerXAnchor.constraint(equalTo: view.centerXAnchor),
      button.topAnchor.constraint(equalTo: imageView.bottomAnchor, constant: 24),
    ])
  }
}

Функция загрузки теперь возвращает stream вместо массива изображений.
Изображения храним в простом массиве, вместо закэшированной таски, как это было ранее.
В рамках загрузочной таски теперь ловим изображения из стрима, параллельно добавляя их в массив. Помимо этого, присваиваем изображение в imageView.image (если его там нет). Причину этого рассмотрим чуть ниже.
Занулям таску в случае ошибки в ней, чтоб была возможность перезапустить ее при следующем нажатии на кнопку.
При нажатии на кнопку пытаемся запустить загрузочную таску и после этого присваиваем в imageView.image случайное изображение из массива (если оно есть). Если массив пуст, это значит, что загрузочная таска еще не успела загрузить ни одного изображения. В таком случае она сама присвоит первое изображение (момент из шага 3).

Нюансы при работе с TaskGroup

В целом принцип работы и предназначение TaskGroup мы разобрали, теперь давайте пробежимся по неочевидным нюансам.

Первый: TaskGroup неявно ожидает завершения всех ChildTask's, даже если вы это не сделаете явно. Посмотрим на примере:

func sleep(_ nanoseconds: UInt64) async {
  try? await Task.sleep(nanoseconds: nanoseconds)
  print("[asyncFunc] finished without errors")
}

func main() async {
  await withTaskGroup(of: Void.self) { group in
    group.addTask { await sleep(1_000_000_000) }
    group.addTask { await sleep(2_000_000_000) }
    group.addTask { await sleep(3_000_000_000) }

    for await _ in group {
      // Завершаем функцию после получения результата из первой ChildTask
      return
    }                                         
  }
  print("[main] finished")
}

Запустив main мы увидим в консоли:

[asyncFunc] finished without errors
[asyncFunc] finished without errors
[asyncFunc] finished without errors
[main] finished

Лог о завершении функции main выведется после завершения всех childTask's, несмотря на то, что мы выходим из функции, не дожидаясь всех child.

Неявное ожидание позволяет масштабировать работу TaskGroup на правила structured concurrency. Каждая Task в structured concurrency перед завершением должна дождаться завершения всех своих детей. То есть дети не должны жить дольше родителя. Это позволяет выстраивать прогнозируемую иерархию из task's.

Дождаться всех дочерних задач из группы можно и явно. Для этого есть метод group.waitForAll(). Он будет полезен в в случаях, когда дочерние задачи ничего не возвращают (то есть просто выполняют какую-то полезную работу), но одновременно с этим после завершения всех нужно еще что-то сделать.

НО. Если мы попробуем провернуть такой же трюк с конструкцией async let, которая тоже в свою очередь создает ChildTask и запускает ее асинхронно, то увидим, что результат будет иным:

func childFunc() async throws {
  try await Task.sleep(nanoseconds: 1_000_000_000) // throws CancellationError
  print("Child task finished")
}

func main() async throws {
  let unstructTask = Task {
    async let childWorkItem = childFunc()
  }

  await unstructTask.value
  print("Main finished")
}

По итогу запуска этого кода увидим в консоли лишь Main finished.

В отличии от TaskGroup, СhildTask, созданная по результату конструкции async let, не ожидается неявно, а просто отменяется. Основное правило при этом не нарушается: child живет не дольше его родителя.

Отличие в поведении основано на предназначении child у TaskGroup и async let. TaskGroup является оберткой над группой child, выполняющих какую-то асинхронную работу, результат которой понадобится целиком (как минимум это предполагается). Поэтому при выходе все задачи должны быть завершены. async let же, в свою очередь, - это просто объявленая асинхронная задача. Если эту задачу никто не ждет, она никому не нужна и она отменяется автоматически.

Следующий нюанс касательно TaskGroup: при выбросе ошибки автоматически отменяются все child Task's. В примере для создания TaskGroup мы пользовались функцией withTaskGroup, но в мире Swift concurrency всегда найдется второй throws вариант. В случае с TaskGroup это withThrowingTaskGroup. Накидаем небольшой пример:

func throwsFunc() async throws {
  throw CancellationError()
}

func nonThrowsFunc() async {
  do {
    try await Task.sleep(nanoseconds: 2_000_000_000)
  } catch {
    print("[nonThrowsFunc] finished with error: \(error.localizedDescription)")
    return
  }

  print("[nonThrowsFunc] finished without errors")
}

func main() async throws {
  try await withThrowingTaskGroup(of: Void.self) { group in
    group.addTask(operation: nonThrowsFunc)
    group.addTask(operation: throwsFunc)

    for try await _ in group { }
    print("[taskGroup] finished")
  }
}

Запустив этот пример, в консоль выведется:

[nonThrowsFunc] finished with error: The operation couldn’t be completed. (Swift.CancellationError error 1.)

В нем у нас есть 2 функции. Одна из них возвращает ошибку, другая - нет. При добавлении обеих функций в группу одна из них сразу выбросит ошибку, а вторая отменится автоматически через TaskGroup. Достаточно интуитивное решение, так как при выбросе ошибки уже не остается смысла в оставшихся подзадачах. Ждать их неявно - waste времени и ресурсов. Не отменять их - waste ресурсов. Самое логичное тут - отмена.

Отменить все активные задачи из группы тоже можно сделать самостоятельно. Для этого существует метод group.cancelAll().

Но есть один момент: все child Task's отменяются только в том случае, если ошибку выбросит именно сама TaskGroup. То есть если мы сделаем так:

func main() async throws {
  try await withThrowingTaskGroup(of: Void.self) { group in
    group.addTask(operation: nonThrowsFunc)
    group.addTask(operation: throwsFunc)
                                                  
    // Не ожидаем завершения подзадач явно (это сделает сама TaskGroup неявно)
    // for try await _ in group { }
    print("[taskGroup] finished")
  }
}

В таком случае наша группа завершится без ошибок (в том числе и nonThrowsFunc) .

И в случае, если мы добавим только две nonThrowsFunc в группу, она может их отменить, если что-то пойдет не так в логике самой группы:

func main() async throws {
  try await withThrowingTaskGroup(of: Void.self) { group in
    // Добавляем 2 nonThrowsFunc
    group.addTask(operation: nonThrowsFunc)
    group.addTask(operation: nonThrowsFunc)

    // Выбрасываем ошибку из тела группы
    throw CancellationError()
    for try await _ in group { }
    print("[taskGroup] finished")
  }
}

В данном случае обе nonThrowsFunc будут отменены.

В таком случае возможно вам захочется воспользоваться конструкцией for try? await in(опциональный try) для игнорирования выброса ошибки. Но сразу же наткнетесь на ошибку компиляции. Данная конструкция невозможна из-за механизма итераторов. Когда вы пишете for try await _ in :

for try await item in sequence { }

Компилятор преобразует это в эквивалент следующего кода:

var iterator = sequence.makeAsyncIterator()
while let item = try await iterator.next() { }

Теперь подставьте вместо обычного try опциональный try?. Появляется неопределенность. Ведь цикл не будет идти дальше после выброса первой же ошибки, хотя, пользуясь for try? await in, интуитивно ожидается, что цикл пройдет по всем элементам и скипнет те, которые возвращают ошибку. Следовательно, если в вашем случае нужно игнорировать ошибку (как это было в нашем примере с изображениями), необходимо уже при выполнении самой child Task скипать ошибку.

Итоги

В данной части мы достаточно плотно поработали с TaskGroup. Разобрали, зачем она нужна и некоторые тонкости работы с ней. Резюмируя, можно сказать, что это очень полезный инструмент для построения групп конкурентных задач, а таких задач достаточно большое количество. На этом наше исследование swift concurrency не заканчивается, впереди еще много интересного, будем на связи.

Полезные ссылки

Swift TaskGroup на примерах