Merge pull request #62946 from tshortli/avoid-defer-in-backdeployed-funcs

[TaskLocals] Avoid use of `defer` in back deployed functions in the standard library

Author: tshortli

stdlib/public/BackDeployConcurrency/TaskCancellation.swift

@@ -46,9 +46,14 @@ public func withTaskCancellationHandler<T>(
   // unconditionally add the cancellation record to the task.
   // if the task was already cancelled, it will be executed right away.
   let record = _taskAddCancellationHandler(handler: handler)
-  defer { _taskRemoveCancellationHandler(record: record) }
-
-  return try await operation()
+  do {
+    let result = try await operation()
+    _taskRemoveCancellationHandler(record: record)
+    return result
+  } catch {
+    _taskRemoveCancellationHandler(record: record)
+    throw error
+  }
 }
 
 @available(SwiftStdlib 5.1, *)

stdlib/public/BackDeployConcurrency/TaskLocal.swift

@@ -145,9 +145,14 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
     _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
 
     _taskLocalValuePush(key: key, value: valueDuringOperation)
-    defer { _taskLocalValuePop() }
-
-    return try await operation()
+    do {
+      let result = try await operation()
+      _taskLocalValuePop()
+      return result
+    } catch {
+      _taskLocalValuePop()
+      throw error
+    }
   }
 
   /// Binds the task-local to the specific value for the duration of the

stdlib/public/Concurrency/TaskCancellation.swift

@@ -46,9 +46,14 @@ public func withTaskCancellationHandler<T>(
   // unconditionally add the cancellation record to the task.
   // if the task was already cancelled, it will be executed right away.
   let record = _taskAddCancellationHandler(handler: handler)
-  defer { _taskRemoveCancellationHandler(record: record) }
-
-  return try await operation()
+  do {
+    let result = try await operation()
+    _taskRemoveCancellationHandler(record: record)
+    return result
+  } catch {
+    _taskRemoveCancellationHandler(record: record)
+    throw error
+  }
 }
 
 @available(SwiftStdlib 5.1, *)

stdlib/public/Concurrency/TaskLocal.swift

@@ -145,9 +145,14 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
     _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
 
     _taskLocalValuePush(key: key, value: valueDuringOperation)
-    defer { _taskLocalValuePop() }
-
-    return try await operation()
+    do {
+      let result = try await operation()
+      _taskLocalValuePop()
+      return result
+    } catch {
+      _taskLocalValuePop()
+      throw error
+    }
   }
 
   /// Binds the task-local to the specific value for the duration of the