Using Swift Concurrency’s task group for tasks with varying output

Earlier, I published a post on Swift Concurrency's task groups. If you haven't read that post yet, and you're not familiar with task groups, I recommend that you read that post first because I won't be explaining task groups in this post. Instead, you will learn about a technique that you can use to work around a limitation of task groups.

Task groups can run a number of child tasks where every child task in the task group produces the same output. This is a hard requirement of the withTaskGroup function. This means that task groups are not the right tool for every job. Sometimes it makes more sense to use async let instead.

In the post where I introduced task groups, I used an example where I needed to fetch a Movie object based on an array of UUIDs. Now let's imagine that our requirements aren't as clear, and we write a function where we receive an array of Descriptor objects that informs us about the type of objects we need to load.

These objects could be either a Movie, or a TVShow. Here's what the Descriptor looks like:

enum MediaType {
    case movie, tvShow
}

struct Descriptor {
    let id: UUID
    let type: MediaType
}

The implementation of Movie and TVShow aren't really relevant in this context. All you need to know is that they can both be loaded from a remote source based on a UUID.

Now let's take a look at the skeleton function that we'll work with:

func fetchMedia(descriptors: [Descriptor]) async -> ???? {
    return await withTaskGroup(of: ????) { group in 
        for descriptor in descriptor {
            group.addTask {
                // do work and return something
            }
        }
    }
}

Notice that I used ???? instead of an actual type for the function's return type and for the type of the task group. We'll need to figure out what we want to return.

One approach would be to create a Media base class and have Movie and TVShow subclass this object. That would work in this case, but it requires us to use classes where we might prefer structs, and it wouldn't work if the the fetched objects weren't so similar.

Instead, we can define an enum and use that as our task output and return type instead. Let's call it a TaskResult:

enum TaskResult {
    case movie(Movie)
    case tvShow(TVShow)
}

Now we can switch on the Descriptor's type, fetch our object, and return a TaskResult where the fetched media is an associated type of our enum case:

func fetchMedia(descriptors: [Descriptor]) async -> [TaskResult] {
    return await withTaskGroup(of: TaskResult.self) { group in 
        for descriptor in descriptor {
            group.addTask {
                switch descriptor.type {
                    case .movie:
                        let movie = await self.fetchMovie(id: descriptor.id)
                        return TaskResult.movie(movie)
                    case .tvShow:
                        let tvShow = await self.fetchShow(id: descriptor.id)
                        return TaskResult.tvShow(tvShow)
                }
            }
        }

        var results = [TaskResult]()

        for await result in group {
            results.append(result)
        }

        return results
    }
}

The nice thing about this approach is that it's easy to scale it into as many types as you need without the need to subclass. That said, I wouldn't recommend this approach in all cases. For example, if you're building a flow similar to the one I show in my post on async let, task groups wouldn't make a lot of sense.

In Summary

Ideally, you only use task groups when all tasks in the group really produce the same output. However, I'm sure there are situations where you need to run an unknown number of tasks based on some input like an array where the tasks don't always produce the same output. In those cases it makes sense to apply the workaround that I've demonstrated in this post.