[dotnet] [bidi] Revisit some core functionality to deserialize without intermediate JsonElement allocation

[nvborisenko]

User description

Core aspects:

• Now we deserialize immediately to "known type" without any intermediate objects
• MessageSuccess<T> which will handle Extensible payload (looking to the future)
• For now just log if we cannot process incoming BiDi message

It is still in "dirty" state, but conceptually it becomes better.

🔗 Related Issues

💥 What does this PR do?

Improve memory allocation. Locating all div nodes on google.com: before 404KB, after 317KB.

🔧 Implementation Notes

💡 Additional Considerations

🔄 Types of changes

• Cleanup (formatting, renaming)
• Bug fix (backwards compatible)
• New feature (non-breaking change which adds functionality and tests!)
• Breaking change (fix or feature that would cause existing functionality to change)

---

PR Type

Enhancement

---

Description

• Refactored BiDi commands to use strongly-typed results, improving deserialization efficiency.

• Introduced EmptyResult as a base class for command results, simplifying result handling.

• Enhanced event handling with a new _eventTypesMap for dynamic event type resolution.

• Removed intermediate JsonElement allocations, optimizing memory usage.

• Replaced int with long for command IDs, ensuring compatibility with larger ranges.

---

Changes walkthrough 📝

	Relevant files
Enhancement	18 files Broker.cs Refactored Broker to handle typed command results and events +119/-36 Command.cs Added result type to Command class for typed results          +11/-3    BiDiJsonSerializerContext.cs Updated serializer context for new result types                    +2/-1      Message.cs Updated Message records for typed results                                +4/-10    CloseCommand.cs Updated CloseCommand to use EmptyResult                                    +1/-1      CreateUserContextCommand.cs Updated CreateUserContextCommand to use typed result          +1/-1      GetClientWindowsCommand.cs Updated GetClientWindowsCommand to use typed result            +2/-2      GetUserContextsCommand.cs Updated GetUserContextsCommand to use typed result              +2/-2      RemoveUserContextCommand.cs Updated RemoveUserContextCommand to use EmptyResult            +1/-1      ActivateCommand.cs Updated ActivateCommand to use EmptyResult                              +1/-1      CaptureScreenshotCommand.cs Updated CaptureScreenshotCommand to use typed result          +2/-2      CloseCommand.cs Updated CloseCommand to use EmptyResult                                    +1/-1      CreateCommand.cs Updated CreateCommand to use typed result                                +2/-2      GetTreeCommand.cs Updated GetTreeCommand to use typed result                              +2/-2      HandleUserPromptCommand.cs Updated HandleUserPromptCommand to use EmptyResult              +1/-1      LocateNodesCommand.cs Updated LocateNodesCommand to use typed result                      +2/-2      NavigateCommand.cs Updated NavigateCommand to use typed result                            +2/-2      PrintCommand.cs Updated PrintCommand to use typed result                                  +2/-2
Cleanup	2 files MessageConverter.cs Removed unused MessageConverter class                                        +0/-45    JsonExtensions.cs Fixed typo in exception message                                                    +2/-1
Additional files	29 files UserContextInfo.cs +3/-1      ReloadCommand.cs +1/-1      SetViewportCommand.cs +1/-1      TraverseHistoryCommand.cs +2/-2      PerformActionsCommand.cs +1/-1      ReleaseActionsCommand.cs +1/-1      SetFilesCommand.cs +1/-1      AddInterceptCommand.cs +2/-2      ContinueRequestCommand.cs +1/-1      ContinueResponseCommand.cs +1/-1      ContinueWithAuthCommand.cs +1/-1      FailRequestCommand.cs +1/-1      ProvideResponseCommand.cs +1/-1      RemoveInterceptCommand.cs +1/-1      SetCacheBehaviorCommand.cs +1/-1      AddPreloadScriptCommand.cs +2/-2      CallFunctionCommand.cs +1/-1      DisownCommand.cs +1/-1      EvaluateCommand.cs +2/-2      GetRealmsCommand.cs +2/-2      RemovePreloadScriptCommand.cs +1/-1      EndCommand.cs +1/-1      NewCommand.cs +2/-2      StatusCommand.cs +2/-2      SubscribeCommand.cs +2/-2      UnsubscribeCommand.cs +2/-2      DeleteCookiesCommand.cs +2/-2      GetCookiesCommand.cs +2/-2      SetCookieCommand.cs +2/-2

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-merge-pro]

PR Reviewer Guide 🔍

(Review updated until commit b560945)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Error Handling The ProcessReceivedMessage method has limited error handling. If _eventTypesMap doesn't contain a requested method key, it will throw an unhandled exception. Consider adding a check before accessing the dictionary. var eventType = _eventTypesMap[method]; var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerContext)!; Potential Race Condition The _eventTypesMap dictionary is accessed from multiple methods without synchronization. Consider using a thread-safe collection like ConcurrentDictionary to prevent potential race conditions. private readonly Dictionary<string, Type> _eventTypesMap = []; Memory Leak If a command fails or times out, the entry in _pendingCommands might not be removed, potentially causing a memory leak. Consider adding cleanup logic for these scenarios. _pendingCommands[command.Id] = (command, tcs);

[qodo-merge-pro]

PR Code Suggestions ✨

Latest suggestions up to b560945
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Missing dictionary key check The code directly accesses _pendingCommands[id.Value] without checking if the command ID exists in the dictionary. If the ID doesn't exist, this will throw a KeyNotFoundException. You should use TryGetValue to safely retrieve the command. dotnet/src/webdriver/BiDi/Communication/Broker.cs [375-384] switch (type) { case "success": if (id is null) throw new JsonException("The remote end responded with 'success' message type, but missed required 'id' property."); - var successCommand = _pendingCommands[id.Value]; + if (!_pendingCommands.TryGetValue(id.Value, out var successCommand)) + { + throw new JsonException($"Received success response for unknown command ID: {id.Value}"); + } + var messageSuccess = JsonSerializer.Deserialize(ref resultReader, successCommand.Item1.ResultType, _jsonSerializerContext)!; successCommand.Item2.SetResult(messageSuccess); _pendingCommands.TryRemove(id.Value, out _); break; Apply this suggestion Suggestion importance[1-10]: 9 __ Why: The suggestion fixes a potential KeyNotFoundException by using TryGetValue instead of direct dictionary access. This is a significant improvement for error handling in a network communication component where unexpected responses may occur.	High
	Missing event method check The code directly accesses _eventTypesMap[method] without checking if the method exists in the dictionary. If the method isn't registered, this will throw a KeyNotFoundException. You should use TryGetValue to safely retrieve the event type or handle unknown events gracefully. dotnet/src/webdriver/BiDi/Communication/Broker.cs [386-395] case "event": if (method is null) throw new JsonException("The remote end responded with 'event' message type, but missed required 'method' property."); - var eventType = _eventTypesMap[method]; + if (!_eventTypesMap.TryGetValue(method, out var eventType)) + { + _logger.Warning($"Received event for unregistered method: {method}"); + break; + } var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerContext)!; var messageEvent = new MessageEvent(method, eventArgs); _pendingEvents.Add(messageEvent); break; Apply this suggestion Suggestion importance[1-10]: 9 __ Why: The suggestion prevents a potential KeyNotFoundException by adding a check before accessing the event type from the dictionary. This is crucial for robust event handling, especially when dealing with potentially unknown event types from remote endpoints.	High
	Null check missing The code is not checking if data is null before creating a new Utf8JsonReader with it. If data is null, this will cause a NullReferenceException when creating the ReadOnlySpan. dotnet/src/webdriver/BiDi/Communication/Broker.cs [320-333] private void ProcessReceivedMessage(byte[]? data) { + if (data == null) + { + throw new ArgumentNullException(nameof(data), "Received data cannot be null"); + } + long? id = default; string? type = default; string? method = default; string? error = default; string? message = default; Utf8JsonReader resultReader = default; Utf8JsonReader paramsReader = default; Utf8JsonReader reader = new(new ReadOnlySpan<byte>(data)); reader.Read(); reader.Read(); // "{" Apply this suggestion Suggestion importance[1-10]: 8 __ Why: The suggestion addresses a potential null reference exception by adding a null check for the 'data' parameter. This is a critical defensive programming practice that prevents runtime crashes when processing network messages.	Medium
Learned best practice	Enhance error message clarity by providing more context and troubleshooting guidance The error message has been corrected from "descriminator" to "discriminator", which is good. However, the error message could be further improved by providing more context about what was being processed and what might have caused the issue. This would help developers troubleshoot the problem more effectively. dotnet/src/webdriver/BiDi/Communication/Json/Internal/JsonExtensions.cs [52] -return discriminator ?? throw new JsonException($"Couldn't determine '{name}' discriminator."); +return discriminator ?? throw new JsonException($"Couldn't determine '{name}' discriminator in the JSON payload. Ensure the property exists and has a valid value."); Apply this suggestion Suggestion importance[1-10]: 6	Low
Update

---

Previous suggestions

✅ Suggestions up to commit 1a4773e

Category	Suggestion	Impact
Possible issue	✅ Check dictionary key exists Suggestion Impact: The commit completely refactored the message processing logic. Instead of directly accessing the dictionary with _pendingCommands[successId], the new implementation uses a different approach with a ProcessReceivedMessage method that handles all message types. While not implementing the exact suggestion, the new code addresses the same issue by using a more comprehensive solution. code diff: + private void ProcessReceivedMessage(byte[]? data) + { + long? id = default; + string? type = default; + string? method = default; + string? error = default; + string? message = default; + Utf8JsonReader resultReader = default; + Utf8JsonReader paramsReader = default; + + Utf8JsonReader reader = new(new ReadOnlySpan<byte>(data)); + reader.Read(); + + reader.Read(); // "{" + + while (reader.TokenType == JsonTokenType.PropertyName) + { + string? propertyName = reader.GetString(); + reader.Read(); + + switch (propertyName) + { + case "id": + id = reader.GetInt64(); + break; + + case "type": + type = reader.GetString(); + break; + + case "method": + method = reader.GetString(); + break; + + case "result": + resultReader = reader; // cloning reader with current position + break; + + case "params": + paramsReader = reader; // cloning reader with current position + break; + + case "error": + error = reader.GetString(); + break; + + case "message": + message = reader.GetString(); + break; + } + + reader.Skip(); + reader.Read(); + } + + switch (type) + { + case "success": + if (id is null) throw new JsonException("The remote end responded with 'success' message type, but missed required 'id' property."); + + var successCommand = _pendingCommands[id.Value]; + var messageSuccess = JsonSerializer.Deserialize(ref resultReader, successCommand.Item1.ResultType, _jsonSerializerContext)!; + successCommand.Item2.SetResult(messageSuccess); + _pendingCommands.TryRemove(id.Value, out _); + break; + + case "event": + if (method is null) throw new JsonException("The remote end responded with 'event' message type, but missed required 'method' property."); + + var eventType = _eventTypesMap[method]; + + var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerContext)!; + + var messageEvent = new MessageEvent(method, eventArgs); + _pendingEvents.Add(messageEvent); + break; + + case "error": + if (id is null) throw new JsonException("The remote end responded with 'error' message type, but missed required 'id' property."); + + var messageError = new MessageError(id.Value) { Error = error, Message = message }; + var errorCommand = _pendingCommands[messageError.Id]; + errorCommand.Item2.SetException(new BiDiException($"{messageError.Error}: {messageError.Message}")); + _pendingCommands.TryRemove(messageError.Id, out _); + break; + } The code doesn't check if the successId exists in the _pendingCommands dictionary before accessing it. This could lead to a KeyNotFoundException if a success message is received for a command ID that is not in the dictionary. dotnet/src/webdriver/BiDi/Communication/Broker.cs [136-142] try { var data = await _transport.ReceiveAsync(cancellationToken).ConfigureAwait(false); Utf8JsonReader utfJsonReader = new(new ReadOnlySpan<byte>(data)); utfJsonReader.Read(); var messageType = utfJsonReader.GetDiscriminator("type"); switch (messageType) { case "success": var successId = int.Parse(utfJsonReader.GetDiscriminator("id")); - var successCommand = _pendingCommands[successId]; - var messageSuccess = JsonSerializer.Deserialize(ref utfJsonReader, successCommand.Item1.ResultType, _jsonSerializerContext); + if (_pendingCommands.TryGetValue(successId, out var successCommand)) + { + var messageSuccess = JsonSerializer.Deserialize(ref utfJsonReader, successCommand.Item1.ResultType, _jsonSerializerContext); - successCommand.Item2.SetResult(messageSuccess); + successCommand.Item2.SetResult(messageSuccess); - _pendingCommands.TryRemove(successId, out _); + _pendingCommands.TryRemove(successId, out _); + } break; Suggestion importance[1-10]: 8 __ Why: This suggestion addresses a potential runtime exception by checking if the key exists in the dictionary before accessing it. Without this check, the code could throw a KeyNotFoundException if a success message is received for a command ID that isn't in the pending commands dictionary.	Medium
	✅ Check event handlers exist Suggestion Impact: The commit completely refactored the message processing logic, including the event handling. The new implementation in ProcessReceivedMessage() uses a different approach that avoids the potential exceptions by using the _eventTypesMap dictionary which is populated during subscription, eliminating the need to access _eventHandlers[method].First().EventArgsType directly. code diff: + private void ProcessReceivedMessage(byte[]? data) + { + long? id = default; + string? type = default; + string? method = default; + string? error = default; + string? message = default; + Utf8JsonReader resultReader = default; + Utf8JsonReader paramsReader = default; + + Utf8JsonReader reader = new(new ReadOnlySpan<byte>(data)); + reader.Read(); + + reader.Read(); // "{" + + while (reader.TokenType == JsonTokenType.PropertyName) + { + string? propertyName = reader.GetString(); + reader.Read(); + + switch (propertyName) + { + case "id": + id = reader.GetInt64(); + break; + + case "type": + type = reader.GetString(); + break; + + case "method": + method = reader.GetString(); + break; + + case "result": + resultReader = reader; // cloning reader with current position + break; + + case "params": + paramsReader = reader; // cloning reader with current position + break; + + case "error": + error = reader.GetString(); + break; + + case "message": + message = reader.GetString(); + break; + } + + reader.Skip(); + reader.Read(); + } + + switch (type) + { + case "success": + var successCommand = _pendingCommands[id.Value]; + var messageSuccess = JsonSerializer.Deserialize(ref resultReader, successCommand.Item1.ResultType, _jsonSerializerContext); + successCommand.Item2.SetResult(messageSuccess); + _pendingCommands.TryRemove(id.Value, out _); + break; + + case "event": + var eventType = _eventTypesMap[method]; + + var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref paramsReader, eventType, _jsonSerializerContext); + + var messageEvent = new MessageEvent(method, eventArgs); + _pendingEvents.Add(messageEvent); + break; + + case "error": + var messageError = new MessageError(id.Value) { Error = error, Message = message }; + var errorCommand = _pendingCommands[messageError.Id]; + errorCommand.Item2.SetException(new BiDiException($"{messageError.Error}: {messageError.Message}")); + _pendingCommands.TryRemove(messageError.Id, out _); + break; + } The code doesn't check if the method exists in the _eventHandlers dictionary or if there are any handlers for that method before accessing it. This could lead to a KeyNotFoundException or InvalidOperationException (when calling First() on an empty collection) if an event is received for a method that has no registered handlers. dotnet/src/webdriver/BiDi/Communication/Broker.cs [145-159] case "event": utfJsonReader.Read(); utfJsonReader.Read(); var method = utfJsonReader.GetString(); utfJsonReader.Read(); - // TODO: Just get type info from existing subscribers, should be better - var type = _eventHandlers[method].First().EventArgsType; + if (_eventHandlers.TryGetValue(method, out var handlers) && handlers.Count > 0) + { + // TODO: Just get type info from existing subscribers, should be better + var type = handlers.First().EventArgsType; - var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref utfJsonReader, type, _jsonSerializerContext); + var eventArgs = (EventArgs)JsonSerializer.Deserialize(ref utfJsonReader, type, _jsonSerializerContext); - var messageEvent = new MessageEvent(method, eventArgs); - _pendingEvents.Add(messageEvent); + var messageEvent = new MessageEvent(method, eventArgs); + _pendingEvents.Add(messageEvent); + } break; [Suggestion has been applied] Suggestion importance[1-10]: 8 __ Why: The suggestion prevents potential KeyNotFoundException and InvalidOperationException by checking if handlers exist for the event method before attempting to access them. This is an important defensive coding practice that improves error handling and robustness.	Medium
Learned best practice	Add null check before accessing properties to prevent potential NullReferenceExceptions The code is accessing args.BiDi without first checking if args is null. This could lead to a NullReferenceException if result.Params returns null. Add a null check before accessing the BiDi property to prevent potential runtime exceptions. dotnet/src/webdriver/BiDi/Communication/Broker.cs [191-193] var args = result.Params; -args.BiDi = _bidi; +if (args != null) +{ + args.BiDi = _bidi; +} Suggestion importance[1-10]: 6	Low

[nvborisenko]

@RenderMichael finally ready for review, actually it simplifies the core functionality. Previously I tried to make remote end Message json polymorphic, but now we deserialize the message manually. It adds complexity (only in 1 place), but it allows us to be performant when deserializing.

Please review changes in Communication directory, ignore other boilerplate changes in Modules directory.

[nvborisenko]

Going forward, truly internal stuff. I need it to make improvements further.

[nvborisenko]

CI fails not related to this PR, merging.

[RenderMichael]

Debug.Assert(reader.TokenType == JsonTokenType.StartObject);, to protect against future refactorings?

Or better yet, maybe we should throw in this case?

[nvborisenko]

Moving fast for now, revisit later as not important for end users.

[RenderMichael]

default: throw new BiDiException($"Unexpected BiDi response: {Encoding.UTF8.GetString(data)}")?

[nvborisenko]

yep, we will introduce handling of "unknown polymorphic types"

[nvborisenko]

Please post an issue to not forget, for now just moving fast to introduce good pattern for BiDi namespace

[RenderMichael]

Not related to this PR: what do you think of renaming this type BiDiCommand? It is more clear and easier for namespacing (we already have a Command type)

[nvborisenko]

No, I don't like any BiDi* type in BiDi namespace. Namespaces are especially created to resolve collisions.

[RenderMichael]

Maybe BaseResult or BiDiResult? Since it is derived by results with values.

[RenderMichael]

Opened #15593

[nvborisenko]

Answered there, I prefer just Result

[RenderMichael]

Empty result is used it mean two things in two different places: base result and void result. The same type serving two different purposes. I propose making this two different types. It will have a real advantage: today, EmptyResult can be assigned a real results value. If we have a separate "no results" type, then it will be guaranteed safe.

[nvborisenko]

We will not have "no results" at all. Any command will return results, read it as public Task<Result> DoSomething().