Add nonisolated `Shared.withLock`

[jshier]

Recently Shared's withLock had @MainActor added due to the need to isolate mutations of the internal value to the main actor. However, composing Shared into other APIs may require locked access to the shared state without such a requirement.

This PR adds an additional overload of withLock that doesn't allow mutation of the stored Value, allowing it to not be isolated to the main actor.

[mbrandonw]

Hi @jshier, thanks for the PR! Can you share a bit of code to demonstrate where this is helpful?

[jshier]

It's generally helpful anywhere you need nonisolated access to the protected Shared state, like any locking mechanism. In this particular instance I was wrapping a @Shared instance in another wrapper to expose additional API. While I can generally guarantee the wrapper is accessed on the main actor, there are specific APIs that should be accessible from outside the main actor, and should be synchronous. In that case, withLock requiring the main actor prevents safe access, while a version of it that doesn't need mutation support and so is not isolated does not.

@MainActor
@propretyWrapper
struct SomeWrapper {
  @Shared(.inMemory("SomeWrapper")) private var state = State()

  var projectedValue: Self { self }
  var wrappedValue: State {
    get { state }
    nonmutating set { state = newValue }
  }

  nonisolated func isSomeState() -> Bool {
    guard let someValue = $state.withLock { $0.someValue } else { return false } // Error

    ...
  }
}

As an aside, I could really use better guidance on this pattern, I really don't know if this is going to work right.

[mbrandonw]

Why is this property wrapper @MainActor? If you dropped that then you could do this right?

guard let someValue = state.someValue else { return false }

[jshier]

Because I use withLock in other places to ensure atomic mutation and it was easiest just to add it. I see no other mechanism to ensure atomic updates.

[jshier]

To step back, I'm writing a wrapper for @Shared that gives access to specific APIs needed for what is essentially a set of locally cached flags. It needs to have both synchronous and asynchronous interfaces to its query and refresh APIs, and eventually needs to respond to a system Notification to reset and refresh its state (this is all integrating with an existing system). It's pretty unclear how this pattern is supposed to work, but here's what I have so far:

@MainActor
@propertyWrapper
struct Flags {
    @Shared(.inMemory("SharedFlags")) private var flags = SharedFeatureFlags()

    var projectedValue: Self { self }

    var wrappedValue: SharedFeatureFlags {
        get { flags }
        nonmutating set { flags = newValue }
    }

    var publisher: AnyPublisher<SharedFeatureFlags, Never> {
        self.$flags.publisher
    }

    nonisolated func refresh(flagSet: FlagSet) {
        Task { await self.refresh(flagSet: flagSet) }
    }

    nonisolated func does(
        _ flag: Flag,
        in set: FlagSet,
        have access: FlagAccess
    ) -> Bool {
        guard let response = $flags.withLock({ $0.responses[flag] }) else { return false }

        if access.contains(.accessible), !response.hasAccessibility(for: flag) {
            return false
        }

        if access.contains(.capability), !response.hasCapability(for: flag) {
            return false
        }

        if access.contains(.subscription), !response.hasSubscription(for: flag) {
            return false
        }

        return true
    }

    func refresh(flagSet: FlagSet) async {
        @Dependency(\.flagClient) var flagsClient

        self.flags.isLoading[group] = true

        var isNewTask = false
        let task: Task<FlagResponse?, Never>
        if let existingTask = flags.refreshTasks[group] {
            task = existingTask
        } else {
            let newTask = Task { await flagsClient.fetchAvailability(flagSet) }
            self.flags.refreshTasks[group] = newTask
            isNewTask = true
            task = newTask
        }

        let response = await task.value

        if isNewTask {
            // Mutate in locks to ensure atomicity.
            self.$flags.withLock { flags in
                flags.refreshTasks[group] = nil
                flags.responses[group] = response
                flags.isLoading[group] = false
            }
        }
    }

    nonisolated func refreshAll() {
        Task { await self.refreshAllTasks() }
    }

    func refreshAll() async {
        let tasks = self.refreshAllTasks()

        for task in tasks {
            await task.value
        }
    }

    @discardableResult
    private func refreshAllTasks() -> [Task<Void, Never>] {
        flags.responses.keys.map { group in
            // Start refreshes in parallel.
            Task { await self.refresh(group: group) }
        }
    }
}

struct SharedFeatureFlags: Equatable {
  var responses: [FlagSet: FlagResponse] = [:]
  var isLoading: [FlagSet: Bool] = [:]
  fileprivate var refreshTasks: [FlagSet: Task<FlagResponse?, Never>] = [:]
}

[mbrandonw]

This example seems like it can drop @MainActor too. There are only two places using withLock. The first is in a synchronous context here:

guard let response = $flags.withLock({ $0.responses[flag] }) else { return false }

But that can just be flags.responses[flag], done without explicit locking (locking implicitly happens in the ValueReference already).

And the second usage of withLock is here:

self.$flags.withLock { flags in
  flags.refreshTasks[group] = nil
  …
}

And this is a legitimate use of withLock, but you are in an async context and so you can do await $flags.withLock.

The only time you would run into trouble is if you had a synchronous endpoint that wanted to mutate shared state. However, that is not a legitimate thing to do, and so you would just need to mark that one single endpoint as @MainActor and then call it from an effect.

[mbrandonw]

Overall, the thing I'm confused about with your PR is that it isn't actually adding anything that isn't already possible with the tools right now. Adding this method:

public func withLock<R>(_ transform: @Sendable (Value) throws -> R) rethrows -> R {
  try transform(self._wrappedValue)
}

…is the exact same as if you were to do transform(self.wrappedValue), and that's all public API. So this API shouldn't be needed, unless I'm misunderstanding something.

[jshier]

This example seems like it can drop @MainActor too. There are only two places using withLock. The first is in a synchronous context here:

guard let response = $flags.withLock({ $0.responses[flag] }) else { return false }

But that can just be flags.responses[flag], done without explicit locking (locking implicitly happens in the ValueReference already).

Perhaps I misunderstand the purpose of withLock, but I don't know of any way to make such direct property access actually safe. That is, even if access to the value is locked, access to that value's properties is not. So while access to flags may be safe, even if that value is ultimately locked, calling .responses[flag] is not within that lock and so unsafe. I've written many versions of withLock that are intended to provide safe access to the value and all of its properties, all the way down, so that's how I used it here.

And the second usage of withLock is here:

self.$flags.withLock { flags in
  flags.refreshTasks[group] = nil
  …
}

And this is a legitimate use of withLock, but you are in an async context and so you can do await $flags.withLock.

This may be safe but it's not equivalent. In addition to the overhead of another suspension point, introducing that second suspension point removes the ordering guarantees I was trying to make, to ensure later callers always have the response available as quickly as possible rather than having to create another Task. If access to the value directly is safe I can just write my own atomic mutator.

Looking at the implementation of withLock again, it doesn't seem to actually do what I expect, which is wrap the work being done under the hood in a lock. Instead, the only lock I see is what lives way down under the currentValue. So how is Shared thread-safe at all?

[jshier]

For that matter, without the @MainActor annotation, how would the wrapper be safe? Regardless of the safety of the stored value, I also need to guarantee ordering and that things are done atomically.