From ad7b31f4f9cf0c041f638cbbffe5d6d052518ed6 Mon Sep 17 00:00:00 2001
From: Caelin Jackson-King <caelin@yelp.com>
Date: Thu, 28 Jun 2018 13:10:43 -0700
Subject: [PATCH] First stab at adding a block to be called and implementing
 didFinishCollectingMetrics on the SessionManager

---
 AFNetworking/AFURLSessionManager.h |  7 +++++++
 AFNetworking/AFURLSessionManager.m | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/AFNetworking/AFURLSessionManager.h b/AFNetworking/AFURLSessionManager.h
index e7ae0d8e1f..72adc12a98 100644
--- a/AFNetworking/AFURLSessionManager.h
+++ b/AFNetworking/AFURLSessionManager.h
@@ -378,6 +378,13 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)setTaskDidCompleteBlock:(nullable void (^)(NSURLSession *session, NSURLSessionTask *task, NSError * _Nullable error))block;
 
+/**
+ Sets a block to be executed when metrics are finalized related to a specific task, as handled by the `NSURLSessionTaskDelegate` method `URLSession:task:didFinishCollectingMetrics:`.
+
+ @param block A block object to be executed when a session task is completed. The block has no return value, and takes three arguments: the session, the task, and any metrics that were collected in the process of executing the task.
+ */
+- (void)setTaskDidFinishCollectingMetricsBlock:(nullable void (^)(NSURLSession *session, NSURLSessionTask *task, NSURLSessionTaskMetrics * _Nullable metrics))block;
+
 ///-------------------------------------------
 /// @name Setting Data Task Delegate Callbacks
 ///-------------------------------------------
diff --git a/AFNetworking/AFURLSessionManager.m b/AFNetworking/AFURLSessionManager.m
index 7a0529b700..35147b894b 100644
--- a/AFNetworking/AFURLSessionManager.m
+++ b/AFNetworking/AFURLSessionManager.m
@@ -97,6 +97,7 @@ static dispatch_group_t url_session_manager_completion_group() {
 typedef NSInputStream * (^AFURLSessionTaskNeedNewBodyStreamBlock)(NSURLSession *session, NSURLSessionTask *task);
 typedef void (^AFURLSessionTaskDidSendBodyDataBlock)(NSURLSession *session, NSURLSessionTask *task, int64_t bytesSent, int64_t totalBytesSent, int64_t totalBytesExpectedToSend);
 typedef void (^AFURLSessionTaskDidCompleteBlock)(NSURLSession *session, NSURLSessionTask *task, NSError *error);
+typedef void (^AFURLSessionTaskDidFinishCollectingMetricsBlock)(NSURLSession *session, NSURLSessionTask *task, NSURLSessionTaskMetrics * metrics);
 
 typedef NSURLSessionResponseDisposition (^AFURLSessionDataTaskDidReceiveResponseBlock)(NSURLSession *session, NSURLSessionDataTask *dataTask, NSURLResponse *response);
 typedef void (^AFURLSessionDataTaskDidBecomeDownloadTaskBlock)(NSURLSession *session, NSURLSessionDataTask *dataTask, NSURLSessionDownloadTask *downloadTask);
@@ -462,6 +463,7 @@ @interface AFURLSessionManager ()
 @property (readwrite, nonatomic, copy) AFURLSessionTaskNeedNewBodyStreamBlock taskNeedNewBodyStream;
 @property (readwrite, nonatomic, copy) AFURLSessionTaskDidSendBodyDataBlock taskDidSendBodyData;
 @property (readwrite, nonatomic, copy) AFURLSessionTaskDidCompleteBlock taskDidComplete;
+@property (readwrite, nonatomic, copy) AFURLSessionTaskDidFinishCollectingMetricsBlock taskDidFinishCollectingMetrics;
 @property (readwrite, nonatomic, copy) AFURLSessionDataTaskDidReceiveResponseBlock dataTaskDidReceiveResponse;
 @property (readwrite, nonatomic, copy) AFURLSessionDataTaskDidBecomeDownloadTaskBlock dataTaskDidBecomeDownloadTask;
 @property (readwrite, nonatomic, copy) AFURLSessionDataTaskDidReceiveDataBlock dataTaskDidReceiveData;
@@ -871,6 +873,10 @@ - (void)setTaskDidCompleteBlock:(void (^)(NSURLSession *session, NSURLSessionTas
     self.taskDidComplete = block;
 }
 
+- (void)setTaskDidFinishCollectingMetricsBlock:(void (^)(NSURLSession * _Nonnull, NSURLSessionTask * _Nonnull, NSURLSessionTaskMetrics * _Nullable))block {
+    self.taskDidFinishCollectingMetrics = block;
+}
+
 #pragma mark -
 
 - (void)setDataTaskDidReceiveResponseBlock:(NSURLSessionResponseDisposition (^)(NSURLSession *session, NSURLSessionDataTask *dataTask, NSURLResponse *response))block {
@@ -1077,6 +1083,22 @@ - (void)URLSession:(NSURLSession *)session
     }
 }
 
+- (void)URLSession:(NSURLSession *)session
+              task:(NSURLSessionTask *)task
+didFinishCollectingMetrics:(NSURLSessionTaskMetrics *)metrics
+{
+    AFURLSessionManagerTaskDelegate *delegate = [self delegateForTask:task];
+
+    // delegate may be nil when completing a task in the background
+    if (delegate && [delegate respondsToSelector:@selector(URLSession:task:didFinishCollectingMetrics:)]) {
+        [delegate URLSession:session task:task didFinishCollectingMetrics:metrics];
+    }
+
+    if (self.taskDidFinishCollectingMetrics) {
+        self.taskDidFinishCollectingMetrics(session, task, metrics);
+    }
+}
+
 #pragma mark - NSURLSessionDataDelegate
 
 - (void)URLSession:(NSURLSession *)session