Added readWithLengthPrefix and writeWithLengthPrefix for thread-safe …

…length-prefixed data messages
OperatorFoundation · Nov 2, 2021 · 3dbea63 · 3dbea63
1 parent 993e6b6
commit 3dbea63
Show file tree

Hide file tree

Showing 2 changed files with 196 additions and 56 deletions.
diff --git a/Package.resolved b/Package.resolved
diff --git a/Sources/Transmission/Connection.swift b/Sources/Transmission/Connection.swift
@@ -87,7 +87,7 @@ public class Connection
     // Reads exactly size bytes
     public func read(size: Int) -> Data?
     {
-        maybeLog(message: "Transmission read called", logger: self.log)
+        maybeLog(message: "Transmission read(size:) called \(Thread.current)", logger: self.log)
         var result: Data?
 
         self.readLock.enter()
@@ -125,7 +125,7 @@ public class Connection
     // reads up to maxSize bytes
     public func read(maxSize: Int) -> Data?
     {
-        maybeLog(message: "Transmission read called", logger: self.log)
+        maybeLog(message: "Transmission read(maxSize:) called \(Thread.current)", logger: self.log)
         var result: Data?
 
         self.readLock.enter()
@@ -159,6 +159,134 @@ public class Connection
 
         return result
     }
+
+    public func readWithLengthPrefix(prefixSizeInBits: Int) -> Data?
+    {
+        maybeLog(message: "Transmission readWithLengthPrefix called \(Thread.current)", logger: self.log)
+        var result: Data?
+
+        self.readLock.enter()
+
+        var maybeCount: Int? = nil
+
+        let countLock = DispatchGroup()
+        countLock.enter()
+        switch prefixSizeInBits
+        {
+            case 8:
+                self.connection.receive(minimumIncompleteLength: 1, maximumLength: 1)
+                {
+                    (maybeData, maybeContext, isComplete, maybeError) in
+
+                    guard maybeError == nil else
+                    {
+                        countLock.leave()
+                        return
+                    }
+
+                    if let data = maybeData
+                    {
+                        if let count = data.maybeNetworkUint8
+                        {
+                             maybeCount = Int(count)
+                        }
+                    }
+                    countLock.leave()
+                }
+            case 16:
+                self.connection.receive(minimumIncompleteLength: 2, maximumLength: 2)
+                {
+                    (maybeData, maybeContext, isComplete, maybeError) in
+
+                    guard maybeError == nil else
+                    {
+                        countLock.leave()
+                        return
+                    }
+
+                    if let data = maybeData
+                    {
+                        if let count = data.maybeNetworkUint16
+                        {
+                            maybeCount = Int(count)
+                        }
+                    }
+                    countLock.leave()
+                }
+            case 32:
+                self.connection.receive(minimumIncompleteLength: 4, maximumLength: 4)
+                {
+                    (maybeData, maybeContext, isComplete, maybeError) in
+
+                    guard maybeError == nil else
+                    {
+                        countLock.leave()
+                        return
+                    }
+
+                    if let data = maybeData
+                    {
+                        if let count = data.maybeNetworkUint32
+                        {
+                            maybeCount = Int(count)
+                        }
+                    }
+                    countLock.leave()
+                }
+            case 64:
+                self.connection.receive(minimumIncompleteLength: 8, maximumLength: 8)
+                {
+                    (maybeData, maybeContext, isComplete, maybeError) in
+
+                    guard maybeError == nil else
+                    {
+                        countLock.leave()
+                        return
+                    }
+
+                    if let data = maybeData
+                    {
+                        if let count = data.maybeNetworkUint64
+                        {
+                            maybeCount = Int(count)
+                        }
+                    }
+                    countLock.leave()
+                }
+            default:
+                countLock.leave()
+        }
+
+        countLock.wait()
+
+        guard let size = maybeCount else
+        {
+            readLock.leave()
+            return nil
+        }
+
+        self.connection.receive(minimumIncompleteLength: size, maximumLength: size)
+        {
+            (maybeData, maybeContext, isComplete, maybeError) in
+
+            guard maybeError == nil else
+            {
+                self.readLock.leave()
+                return
+            }
+
+            if let data = maybeData
+            {
+                result = data
+            }
+
+            self.readLock.leave()
+        }
+
+        readLock.wait()
+
+        return result
+    }
 
     public func write(string: String) -> Bool
     {
@@ -168,7 +296,7 @@ public class Connection
 
     public func write(data: Data) -> Bool
     {
-        maybeLog(message: "Transmission write called", logger: self.log)
+        maybeLog(message: "Transmission write called \(Thread.current)", logger: self.log)
         var success = false
 
         self.writeLock.enter()
@@ -190,10 +318,74 @@ public class Connection
 
         self.writeLock.wait()
 
-        maybeLog(message: "Transmission write finished", logger: self.log)
+        maybeLog(message: "Transmission write finished \(Thread.current)", logger: self.log)
 
         return success
     }
+
+    public func writeWithLengthPrefix(data: Data, prefixSizeInBits: Int) -> Bool
+    {
+        var maybeCountData: Data? = nil
+
+        switch prefixSizeInBits
+        {
+            case 8:
+                let count = UInt8(data.count)
+                maybeCountData = count.maybeNetworkData
+            case 16:
+                let count = UInt16(data.count)
+                maybeCountData = count.maybeNetworkData
+            case 32:
+                let count = UInt32(data.count)
+                maybeCountData = count.maybeNetworkData
+            case 64:
+                let count = UInt64(data.count)
+                maybeCountData = count.maybeNetworkData
+            default:
+                return false
+        }
+
+        guard let countData = maybeCountData else {return false}
+
+        maybeLog(message: "Transmission writeWithLengthPrefix called \(Thread.current)", logger: self.log)
+        var success = false
+
+        self.writeLock.enter()
+
+        self.connection.send(content: countData, contentContext: NWConnection.ContentContext.defaultMessage, isComplete: false, completion: NWConnection.SendCompletion.contentProcessed(
+            {
+                (maybeError) in
+
+                guard maybeError == nil else
+                {
+                    success = false
+                    self.writeLock.leave()
+                    return
+                }
+
+                self.connection.send(content: data, contentContext: NWConnection.ContentContext.defaultMessage, isComplete: false, completion: NWConnection.SendCompletion.contentProcessed(
+                    {
+                        (maybeError) in
+
+                        guard maybeError == nil else
+                        {
+                            success = false
+                            self.writeLock.leave()
+                            return
+                        }
+
+                        success = true
+                        self.writeLock.leave()
+                        return
+                    }))
+            }))
+
+        self.writeLock.wait()
+
+        maybeLog(message: "Transmission writeWithLengthPrefix finished \(Thread.current)", logger: self.log)
+
+        return success
+    }
 }
 
 public enum ConnectionType