Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix stream cancelation cancels underlying stream #395

Merged
merged 5 commits into from
Sep 3, 2024
+246 −103
Merged

Fix stream cancelation cancels underlying stream #395

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
687e31d
update the stream to correctly terminate the underlying stream when c…
nplasterer Sep 2, 2024
dc4bbb3
fix up a bunch of test warnings
nplasterer Sep 2, 2024
de02285
update group streaming as well
nplasterer Sep 2, 2024
e952444
bump the pod
nplasterer Sep 2, 2024
4698c19
make sure all streams cancel correctly
nplasterer Sep 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 96 additions & 34 deletions Sources/XMTPiOS/Conversations.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,29 +130,52 @@ public actor Conversations {

public func streamGroups() async throws -> AsyncThrowingStream<Group, Error> {
AsyncThrowingStream { continuation in
Task {
let task = Task {
let groupCallback = GroupStreamCallback(client: self.client) { group in
guard !Task.isCancelled else {
continuation.finish()
return
}
continuation.yield(group)
}
guard let stream = try await self.client.v3Client?.conversations().stream(callback: groupCallback) else {
guard let stream = await self.client.v3Client?.conversations().stream(callback: groupCallback) else {
continuation.finish(throwing: GroupError.streamingFailure)
return
}

self.streamHolder.stream = stream
continuation.onTermination = { @Sendable reason in
stream.end()
}
}

continuation.onTermination = { @Sendable reason in
task.cancel()
self.streamHolder.stream?.end()
}
}
}

private func streamGroupConversations() -> AsyncThrowingStream<Conversation, Error> {
AsyncThrowingStream { continuation in
Task {
self.streamHolder.stream = try await self.client.v3Client?.conversations().stream(
let task = Task {
self.streamHolder.stream = await self.client.v3Client?.conversations().stream(
callback: GroupStreamCallback(client: self.client) { group in
guard !Task.isCancelled else {
continuation.finish()
return
}
continuation.yield(Conversation.group(group))
}
)
continuation.onTermination = { @Sendable reason in
self.streamHolder.stream?.end()
}
}

continuation.onTermination = { @Sendable reason in
task.cancel()
self.streamHolder.stream?.end()
}
}
}
Expand Down Expand Up @@ -383,63 +406,88 @@ public actor Conversations {

public func streamAllGroupMessages() -> AsyncThrowingStream<DecodedMessage, Error> {
AsyncThrowingStream { continuation in
Task {
let messageCallback = MessageCallback(client: self.client) { message in
if let decodedMessage = MessageV3(client: self.client, ffiMessage: message).decodeOrNull() {
continuation.yield(decodedMessage)
let task = Task {
self.streamHolder.stream = await self.client.v3Client?.conversations().streamAllMessages(
messageCallback: MessageCallback(client: self.client) { message in
guard !Task.isCancelled else {
continuation.finish()
self.streamHolder.stream?.end() // End the stream upon cancellation
return
}
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decode())
} catch {
print("Error onMessage \(error)")
}
}
}
guard let stream = try await client.v3Client?.conversations().streamAllMessages(messageCallback: messageCallback) else {
continuation.finish(throwing: GroupError.streamingFailure)
return
}
continuation.onTermination = { @Sendable reason in
stream.end()
}
)
}

continuation.onTermination = { _ in
task.cancel()
self.streamHolder.stream?.end()
}
}
}

public func streamAllMessages(includeGroups: Bool = false) async throws -> AsyncThrowingStream<DecodedMessage, Error> {
public func streamAllMessages(includeGroups: Bool = false) -> AsyncThrowingStream<DecodedMessage, Error> {
AsyncThrowingStream<DecodedMessage, Error> { continuation in
@Sendable func forwardStreamToMerged(stream: AsyncThrowingStream<DecodedMessage, Error>) async {
do {
var iterator = stream.makeAsyncIterator()
while let element = try await iterator.next() {
guard !Task.isCancelled else {
continuation.finish()
self.streamHolder.stream?.end()
return
}
continuation.yield(element)
}
continuation.finish()
} catch {
continuation.finish(throwing: error)
}
}
Task {

let task = Task {
await forwardStreamToMerged(stream: streamAllV2Messages())
}

if includeGroups {
Task {
await forwardStreamToMerged(stream: streamAllGroupMessages())
}
}

continuation.onTermination = { _ in
task.cancel()
self.streamHolder.stream?.end()
}
}
}

public func streamAllGroupDecryptedMessages() -> AsyncThrowingStream<DecryptedMessage, Error> {
AsyncThrowingStream { continuation in
Task {
do {
self.streamHolder.stream = try await self.client.v3Client?.conversations().streamAllMessages(
messageCallback: MessageCallback(client: self.client) { message in
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decrypt())
} catch {
print("Error onMessage \(error)")
}
let task = Task {
self.streamHolder.stream = await self.client.v3Client?.conversations().streamAllMessages(
messageCallback: MessageCallback(client: self.client) { message in
guard !Task.isCancelled else {
continuation.finish()
self.streamHolder.stream?.end() // End the stream upon cancellation
return
}
)
} catch {
print("STREAM ERR: \(error)")
}
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decrypt())
} catch {
print("Error onMessage \(error)")
}
}
)
}

continuation.onTermination = { _ in
task.cancel()
self.streamHolder.stream?.end()
}
}
}
Expand All @@ -450,23 +498,37 @@ public actor Conversations {
do {
var iterator = stream.makeAsyncIterator()
while let element = try await iterator.next() {
guard !Task.isCancelled else {
continuation.finish()
self.streamHolder.stream?.end()
return
}
continuation.yield(element)
}
continuation.finish()
} catch {
continuation.finish(throwing: error)
}
}
Task {
await forwardStreamToMerged(stream: try streamAllV2DecryptedMessages())

let task = Task {
await forwardStreamToMerged(stream: streamAllV2DecryptedMessages())
}
if (includeGroups) {

if includeGroups {
Task {
await forwardStreamToMerged(stream: streamAllGroupDecryptedMessages())
}
}

continuation.onTermination = { _ in
task.cancel()
self.streamHolder.stream?.end()
}
}
}




func streamAllV2DecryptedMessages() -> AsyncThrowingStream<DecryptedMessage, Error> {
Expand Down
68 changes: 44 additions & 24 deletions Sources/XMTPiOS/Group.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,41 +293,61 @@ public struct Group: Identifiable, Equatable, Hashable {

public func streamMessages() -> AsyncThrowingStream<DecodedMessage, Error> {
AsyncThrowingStream { continuation in
Task.detached {
do {
self.streamHolder.stream = try await ffiGroup.stream(
messageCallback: MessageCallback(client: self.client) { message in
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decode())
} catch {
print("Error onMessage \(error)")
}
let task = Task.detached {
self.streamHolder.stream = await self.ffiGroup.stream(
messageCallback: MessageCallback(client: self.client) { message in
guard !Task.isCancelled else {
continuation.finish()
return
}
)
} catch {
print("STREAM ERR: \(error)")
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decode())
} catch {
print("Error onMessage \(error)")
continuation.finish(throwing: error)
}
}
)

continuation.onTermination = { @Sendable reason in
self.streamHolder.stream?.end()
}
}

continuation.onTermination = { @Sendable reason in
task.cancel()
self.streamHolder.stream?.end()
}
}
}

public func streamDecryptedMessages() -> AsyncThrowingStream<DecryptedMessage, Error> {
AsyncThrowingStream { continuation in
Task.detached {
do {
self.streamHolder.stream = try await ffiGroup.stream(
messageCallback: MessageCallback(client: self.client) { message in
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decrypt())
} catch {
print("Error onMessage \(error)")
}
let task = Task.detached {
self.streamHolder.stream = await self.ffiGroup.stream(
messageCallback: MessageCallback(client: self.client) { message in
guard !Task.isCancelled else {
continuation.finish()
return
}
)
} catch {
print("STREAM ERR: \(error)")
do {
continuation.yield(try MessageV3(client: self.client, ffiMessage: message).decrypt())
} catch {
print("Error onMessage \(error)")
continuation.finish(throwing: error)
}
}
)

continuation.onTermination = { @Sendable reason in
self.streamHolder.stream?.end()
}
}

continuation.onTermination = { @Sendable reason in
task.cancel()
self.streamHolder.stream?.end()
}
}
}

Expand Down
18 changes: 9 additions & 9 deletions Tests/XMTPTests/ConversationTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,16 +84,16 @@ class ConversationTests: XCTestCase {
}

func testDoesNotAllowConversationWithSelf() async throws {
let expectation = expectation(description: "convo with self throws")
let expectation = XCTestExpectation(description: "convo with self throws")
let client = aliceClient!

do {
try await client.conversations.newConversation(with: alice.walletAddress)
_ = try await client.conversations.newConversation(with: alice.walletAddress)
} catch {
expectation.fulfill()
}

wait(for: [expectation], timeout: 0.1)
await fulfillment(of: [expectation], timeout: 3)
}

func testCanStreamConversationsV2() async throws {
Expand All @@ -103,7 +103,7 @@ class ConversationTests: XCTestCase {

let wallet2 = try PrivateKey.generate()
let client2 = try await Client.create(account: wallet2, options: options)
let expectation1 = expectation(description: "got a conversation")
let expectation1 = XCTestExpectation(description: "got a conversation")
expectation1.expectedFulfillmentCount = 2

Task(priority: .userInitiated) {
Expand Down Expand Up @@ -140,7 +140,7 @@ class ConversationTests: XCTestCase {

try await conversation2.send(content: "hi from new wallet")

await waitForExpectations(timeout: 30)
await fulfillment(of: [expectation1], timeout: 30)
}

func publishLegacyContact(client: Client) async throws {
Expand All @@ -161,7 +161,7 @@ class ConversationTests: XCTestCase {
return
}

let expectation = expectation(description: "got a message")
let expectation = XCTestExpectation(description: "got a message")

Task(priority: .userInitiated) {
for try await message in conversation.streamMessages() {
Expand All @@ -174,7 +174,7 @@ class ConversationTests: XCTestCase {
// Stream a message
try await conversation.send(content: "hi alice")

await waitForExpectations(timeout: 3)
await fulfillment(of: [expectation], timeout: 3)
}

func testCanLoadV2Messages() async throws {
Expand Down Expand Up @@ -458,7 +458,7 @@ class ConversationTests: XCTestCase {
XCTAssertTrue(isAllowed)

try await bobClient.contacts.deny(addresses: [alice.address])
try await bobClient.contacts.refreshConsentList()
_ = try await bobClient.contacts.refreshConsentList()

let isDenied = (try await bobConversation.consentState()) == .denied

Expand Down Expand Up @@ -491,7 +491,7 @@ class ConversationTests: XCTestCase {
XCTAssertTrue(isUnknown)

try await aliceConversation.send(content: "hey bob")
try await aliceClient.contacts.refreshConsentList()
_ = try await aliceClient.contacts.refreshConsentList()
let isNowAllowed = (try await aliceConversation.consentState()) == .allowed

// Conversations you send a message to get marked as allowed
Expand Down
Loading
Loading