fix: Make Context thread-safe (#853)

Author: jbelkins

Diff for: Sources/ClientRuntime/Idempotency/Context+Idempotency.swift

@@ -12,7 +12,7 @@ import class Smithy.ContextBuilder
 extension Context {
 
     public func getIdempotencyTokenGenerator() -> IdempotencyTokenGenerator {
-        return attributes.get(key: idempotencyTokenGeneratorKey)!
+        get(key: idempotencyTokenGeneratorKey)!
     }
 }
 

Diff for: Sources/Smithy/Context.swift

@@ -5,46 +5,63 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+import class Foundation.NSRecursiveLock
+
 public class Context {
-    public var attributes: Attributes
+    private var _attributes: Attributes
+    private let lock = NSRecursiveLock()
 
     public init(attributes: Attributes) {
-        self.attributes = attributes
+        self._attributes = attributes
     }
 
     public func toBuilder() -> ContextBuilder {
-        let builder = ContextBuilder()
-        builder.attributes = self.attributes
-        return builder
+        ContextBuilder(attributes: accessAttributes())
     }
 
     public func getLogger() -> LogAgent? {
-        return attributes.get(key: AttributeKeys.logger)
+        return accessAttributes().get(key: AttributeKeys.logger)
+    }
+
+    @discardableResult
+    private func accessAttributes(accessor: ((inout Attributes) -> Void)? = nil) -> Attributes {
+        lock.lock()
+        defer { lock.unlock() }
+        accessor?(&_attributes)
+        return _attributes
     }
 }
 
 extension Context {
+
     public func get<T>(key: AttributeKey<T>) -> T? {
-        self.attributes.get(key: key)
+        accessAttributes().get(key: key)
     }
 
     public func contains<T>(key: AttributeKey<T>) -> Bool {
-        self.attributes.contains(key: key)
+        accessAttributes().contains(key: key)
     }
 
     public func set<T>(key: AttributeKey<T>, value: T?) {
-        self.attributes.set(key: key, value: value)
+        accessAttributes { attributes in
+            attributes.set(key: key, value: value)
+        }
     }
 
     public func remove<T>(key: AttributeKey<T>) {
-        self.attributes.remove(key: key)
+        accessAttributes { attributes in
+            attributes.remove(key: key)
+        }
     }
 }
 
 public class ContextBuilder {
-    public init() {}
 
-    public var attributes: Attributes = Attributes()
+    public init(attributes: Attributes = Attributes()) {
+        self.attributes = attributes
+    }
+
+    public var attributes: Attributes
 
     // We follow the convention of returning the builder object
     // itself from any configuration methods, and by adding the

Diff for: Sources/SmithyChecksumsAPI/Context+Checksum.swift

@@ -11,8 +11,8 @@ import struct Smithy.AttributeKey
 public extension Context {
 
     var checksum: ChecksumAlgorithm? {
-        get { attributes.get(key: checksumKey) }
-        set { attributes.set(key: checksumKey, value: newValue) }
+        get { get(key: checksumKey) }
+        set { set(key: checksumKey, value: newValue) }
     }
 
     var checksumString: String? { self.checksum?.toString() }

Diff for: Sources/SmithyEventStreamsAPI/Context+EventStreamsAPI.swift

@@ -11,8 +11,8 @@ import struct Smithy.AttributeKey
 public extension Context {
 
     var messageEncoder: MessageEncoder? {
-        get { attributes.get(key: messageEncoderKey) }
-        set { attributes.set(key: messageEncoderKey, value: newValue) }
+        get { get(key: messageEncoderKey) }
+        set { set(key: messageEncoderKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyEventStreamsAuthAPI/Context+EventStreamsAuthAPI.swift

@@ -11,8 +11,8 @@ import struct Smithy.AttributeKey
 public extension Context {
 
     var messageSigner: MessageSigner? {
-        get { attributes.get(key: messageSignerKey) }
-        set { attributes.set(key: messageSignerKey, value: newValue) }
+        get { get(key: messageSignerKey) }
+        set { set(key: messageSignerKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyHTTPAPI/Context+HTTP.swift

@@ -15,73 +15,73 @@ import struct Foundation.TimeInterval
 extension Context {
 
     public var httpResponse: HTTPResponse? {
-        get { attributes.get(key: httpResponseKey) }
-        set { attributes.set(key: httpResponseKey, value: newValue) }
+        get { get(key: httpResponseKey) }
+        set { set(key: httpResponseKey, value: newValue) }
     }
 
     public var expiration: TimeInterval {
-        get { attributes.get(key: expirationKey) ?? 0 }
-        set { attributes.set(key: expirationKey, value: newValue) }
+        get { get(key: expirationKey) ?? 0 }
+        set { set(key: expirationKey, value: newValue) }
     }
 
     public var host: String? {
-        get { attributes.get(key: hostKey) }
-        set { attributes.set(key: hostKey, value: newValue) }
+        get { get(key: hostKey) }
+        set { set(key: hostKey, value: newValue) }
     }
 
     public var hostPrefix: String? {
-        get { attributes.get(key: hostPrefixKey) }
-        set { attributes.set(key: hostPrefixKey, value: newValue) }
+        get { get(key: hostPrefixKey) }
+        set { set(key: hostPrefixKey, value: newValue) }
     }
 
     public var method: HTTPMethodType {
-        get { attributes.get(key: methodKey) ?? .get }
-        set { attributes.set(key: methodKey, value: newValue) }
+        get { get(key: methodKey) ?? .get }
+        set { set(key: methodKey, value: newValue) }
     }
 
     public func getOperation() -> String? {
-        return attributes.get(key: SmithyHTTPAPIKeys.operation)
+        get(key: SmithyHTTPAPIKeys.operation)
     }
 
     /// The partition ID to be used for this context.
     ///
     /// Requests made with the same partition ID will be grouped together for retry throttling purposes.
     /// If no partition ID is provided, requests will be partitioned based on the hostname.
     public var partitionID: String? {
-        get { attributes.get(key: partitionIDKey) }
-        set { attributes.set(key: partitionIDKey, value: newValue) }
+        get { get(key: partitionIDKey) }
+        set { set(key: partitionIDKey, value: newValue) }
     }
 
     public var path: String {
-        get { attributes.get(key: pathKey)! }
-        set { attributes.set(key: pathKey, value: newValue) }
+        get { get(key: pathKey)! }
+        set { set(key: pathKey, value: newValue) }
     }
 
     public func getRegion() -> String? {
-        return attributes.get(key: SmithyHTTPAPIKeys.region)
+        return get(key: SmithyHTTPAPIKeys.region)
     }
 
     public func getServiceName() -> String {
-        return attributes.get(key: SmithyHTTPAPIKeys.serviceName)!
+        return get(key: SmithyHTTPAPIKeys.serviceName)!
     }
 
     public var signingName: String? {
-        get { attributes.get(key: signingNameKey) }
-        set { attributes.set(key: signingNameKey, value: newValue) }
+        get { get(key: signingNameKey) }
+        set { set(key: signingNameKey, value: newValue) }
     }
 
     public var signingRegion: String? {
-        get { attributes.get(key: signingRegionKey) }
-        set { attributes.set(key: signingRegionKey, value: newValue) }
+        get { get(key: signingRegionKey) }
+        set { set(key: signingRegionKey, value: newValue) }
     }
 
     public func hasUnsignedPayloadTrait() -> Bool {
-        return attributes.get(key: SmithyHTTPAPIKeys.hasUnsignedPayloadTrait) ?? false
+        return get(key: SmithyHTTPAPIKeys.hasUnsignedPayloadTrait) ?? false
     }
 
     public var isBidirectionalStreamingEnabled: Bool {
-        get { attributes.get(key: isBidirectionalStreamingEnabledKey) ?? false }
-        set { attributes.set(key: isBidirectionalStreamingEnabledKey, value: newValue) }
+        get { get(key: isBidirectionalStreamingEnabledKey) ?? false }
+        set { set(key: isBidirectionalStreamingEnabledKey, value: newValue) }
     }
 
     /// Returns `true` if the request should use `http2` and only `http2` without falling back to `http1`

Diff for: Sources/SmithyHTTPAuthAPI/Context/Context+Chunked.swift

@@ -10,8 +10,8 @@ import struct Smithy.AttributeKey
 
 public extension Context {
     var isChunkedEligibleStream: Bool? {
-        get { attributes.get(key: isChunkedEligibleStreamKey) }
-        set { attributes.set(key: isChunkedEligibleStreamKey, value: newValue) }
+        get { get(key: isChunkedEligibleStreamKey) }
+        set { set(key: isChunkedEligibleStreamKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyHTTPAuthAPI/Context/Context+EstimatedSkew.swift

@@ -12,13 +12,13 @@ import struct Smithy.AttributeKey
 
 public extension Context {
     var estimatedSkew: TimeInterval? {
-        get { attributes.get(key: estimatedSkewKey) }
-        set { attributes.set(key: estimatedSkewKey, value: newValue) }
+        get { get(key: estimatedSkewKey) }
+        set { set(key: estimatedSkewKey, value: newValue) }
     }
 
     var socketTimeout: TimeInterval? {
-        get { attributes.get(key: socketTimeoutKey) }
-        set { attributes.set(key: socketTimeoutKey, value: newValue) }
+        get { get(key: socketTimeoutKey) }
+        set { set(key: socketTimeoutKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyHTTPAuthAPI/Context/Context+HTTPAuth.swift

@@ -11,26 +11,27 @@ import struct Smithy.Attributes
 import struct Smithy.AttributeKey
 
 public extension Context {
+
     func getAuthSchemes() -> Attributes? {
-        return attributes.get(key: authSchemesKey)
+        get(key: authSchemesKey)
     }
 
     var selectedAuthScheme: SelectedAuthScheme? {
-        get { attributes.get(key: selectedAuthSchemeKey) }
-        set { attributes.set(key: selectedAuthSchemeKey, value: newValue) }
+        get { get(key: selectedAuthSchemeKey) }
+        set { set(key: selectedAuthSchemeKey, value: newValue) }
     }
 
     func setSelectedAuthScheme(_ value: SelectedAuthScheme?) {
-        attributes.set(key: selectedAuthSchemeKey, value: value)
+        set(key: selectedAuthSchemeKey, value: value)
     }
 
     func getAuthSchemeResolver() -> AuthSchemeResolver? {
-        return attributes.get(key: authSchemeResolverKey)
+        get(key: authSchemeResolverKey)
     }
 
     var signingAlgorithm: SigningAlgorithm? {
-        get { attributes.get(key: signingAlgorithmKey) }
-        set { attributes.set(key: signingAlgorithmKey, value: newValue) }
+        get { get(key: signingAlgorithmKey) }
+        set { set(key: signingAlgorithmKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyHTTPAuthAPI/Context/Context+RequestSignature.swift

@@ -12,8 +12,8 @@ import struct Smithy.AttributeKey
 public extension Context {
 
     var requestSignature: String {
-        get { attributes.get(key: requestSignatureKey) ?? "" }
-        set { attributes.set(key: requestSignatureKey, value: newValue) }
+        get { get(key: requestSignatureKey) ?? "" }
+        set { set(key: requestSignatureKey, value: newValue) }
     }
 }
 

Diff for: Sources/SmithyIdentityAPI/Context/Context+FlowType.swift

@@ -13,7 +13,7 @@ public let flowTypeKey = AttributeKey<FlowType>(name: "FlowType")
 
 extension Context {
     public func getFlowType() -> FlowType {
-        return attributes.get(key: flowTypeKey) ?? .NORMAL
+        get(key: flowTypeKey) ?? .NORMAL
     }
 }
 

Diff for: Sources/SmithyIdentityAPI/Context/Context+IdentityResolver.swift

@@ -12,7 +12,7 @@ import class Smithy.ContextBuilder
 
 extension Context {
     public func getIdentityResolvers() -> Attributes? {
-        return attributes.get(key: identityResolversKey)
+        get(key: identityResolversKey)
     }
 }
 

Diff for: Tests/ClientRuntimeTests/InterceptorTests/InterceptorTests.swift

@@ -123,7 +123,7 @@ class InterceptorTests: XCTestCase {
         let updatedInput = interceptorContext.getInput()
         XCTAssertEqual(updatedInput.property, "bar")
         XCTAssertEqual(updatedInput.otherProperty, 1)
-        XCTAssertEqual(interceptorContext.getAttributes().attributes.get(key: AttributeKey(name: "foo")), "bar")
+        XCTAssertEqual(interceptorContext.getAttributes().get(key: AttributeKey(name: "foo")), "bar")
         XCTAssertEqual(interceptorContext.getRequest().headers.value(for: "foo"), "bar")
         XCTAssertEqual(interceptorContext.getRequest().headers.value(for: "otherProperty"), "1")
     }