Skip to content

Commit

Permalink
[releases/1.2] service: Fix tag deserialization for large messages (#554
Browse files Browse the repository at this point in the history
)

service: Fix tag deserialization for large messages (#552)

* tests: Add test cases for big messages (100 fields)

* service: Fix tag deserialization for large messages

Tags are i32 varints, but we were only deserializing the first byte.

Simplify some related details:
- Decoding a varint returns the updated position, so we don't need to use encoder._TagSize()
- We can directly construct a memory view once instead of constructing a BytesIO for each loop iteration.

* tests: Use pytest.approx instead of choosing exactly comparable floats
  • Loading branch information
bkeryan authored Dec 13, 2023
1 parent 405f063 commit 67f8b2a
Show file tree
Hide file tree
Showing 7 changed files with 534 additions and 24 deletions.
31 changes: 7 additions & 24 deletions ni_measurementlink_service/_internal/parameter/serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
from typing import Any, Dict, Sequence, cast

from google.protobuf.descriptor import FieldDescriptor
from google.protobuf.internal import encoder
from google.protobuf.internal.decoder import ( # type: ignore[attr-defined]
_DecodeSignedVarint32,
)
from google.protobuf.message import Message

from ni_measurementlink_service._internal.parameter import serialization_strategy
Expand Down Expand Up @@ -110,25 +112,6 @@ def serialize_default_values(parameter_metadata_dict: Dict[int, ParameterMetadat
return serialize_parameters(parameter_metadata_dict, default_value_parameter_array)


def _get_field_index(parameter_bytes: bytes, tag_position: int) -> int:
"""Get the Filed Index based on the tag's position.
The tag Position should be the index of the TagValue in the ByteArray for valid field index.
Args
----
parameter_bytes (bytes): Serialized bytes
tag_position (int): Tag position
Returns
-------
int: Filed index of the Tag Position
"""
return parameter_bytes[tag_position] >> _GRPC_WIRE_TYPE_BIT_WIDTH


def _get_overlapping_parameters(
parameter_metadata_dict: Dict[int, ParameterMetadata], parameter_bytes: bytes
) -> Dict[int, Any]:
Expand All @@ -152,8 +135,10 @@ def _get_overlapping_parameters(
# inner_decoder update the overlapping_parameters
overlapping_parameters_by_id: Dict[int, Any] = {}
position = 0
parameter_bytes_memory_view = memoryview(parameter_bytes)
while position < len(parameter_bytes):
field_index = _get_field_index(parameter_bytes, position)
(tag, position) = _DecodeSignedVarint32(parameter_bytes_memory_view, position)
field_index = tag >> _GRPC_WIRE_TYPE_BIT_WIDTH
if field_index not in parameter_metadata_dict:
raise Exception(
f"Error occurred while reading the parameter - given protobuf index '{field_index}' is invalid."
Expand All @@ -163,11 +148,9 @@ def _get_overlapping_parameters(
field_metadata.type, field_metadata.repeated, field_metadata.message_type
)
inner_decoder = decoder(field_index, cast(FieldDescriptor, field_index))
parameter_bytes_io = BytesIO(parameter_bytes)
parameter_bytes_memory_view = parameter_bytes_io.getbuffer()
position = inner_decoder(
parameter_bytes_memory_view,
position + encoder._TagSize(field_index), # type: ignore[attr-defined]
position,
len(parameter_bytes),
cast(Message, None), # unused - See serialization_strategy._vector_decoder._new_default
cast(Dict[FieldDescriptor, Any], overlapping_parameters_by_id),
Expand Down
105 changes: 105 additions & 0 deletions tests/assets/bigmessage.proto
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
syntax = "proto3";
package ni.measurementlink.measurement.tests.v1;

message BigMessage {
double field1 = 1;
double field2 = 2;
double field3 = 3;
double field4 = 4;
double field5 = 5;
double field6 = 6;
double field7 = 7;
double field8 = 8;
double field9 = 9;
double field10 = 10;
double field11 = 11;
double field12 = 12;
double field13 = 13;
double field14 = 14;
double field15 = 15;
double field16 = 16;
double field17 = 17;
double field18 = 18;
double field19 = 19;
double field20 = 20;
double field21 = 21;
double field22 = 22;
double field23 = 23;
double field24 = 24;
double field25 = 25;
double field26 = 26;
double field27 = 27;
double field28 = 28;
double field29 = 29;
double field30 = 30;
double field31 = 31;
double field32 = 32;
double field33 = 33;
double field34 = 34;
double field35 = 35;
double field36 = 36;
double field37 = 37;
double field38 = 38;
double field39 = 39;
double field40 = 40;
double field41 = 41;
double field42 = 42;
double field43 = 43;
double field44 = 44;
double field45 = 45;
double field46 = 46;
double field47 = 47;
double field48 = 48;
double field49 = 49;
double field50 = 50;
double field51 = 51;
double field52 = 52;
double field53 = 53;
double field54 = 54;
double field55 = 55;
double field56 = 56;
double field57 = 57;
double field58 = 58;
double field59 = 59;
double field60 = 60;
double field61 = 61;
double field62 = 62;
double field63 = 63;
double field64 = 64;
double field65 = 65;
double field66 = 66;
double field67 = 67;
double field68 = 68;
double field69 = 69;
double field70 = 70;
double field71 = 71;
double field72 = 72;
double field73 = 73;
double field74 = 74;
double field75 = 75;
double field76 = 76;
double field77 = 77;
double field78 = 78;
double field79 = 79;
double field80 = 80;
double field81 = 81;
double field82 = 82;
double field83 = 83;
double field84 = 84;
double field85 = 85;
double field86 = 86;
double field87 = 87;
double field88 = 88;
double field89 = 89;
double field90 = 90;
double field91 = 91;
double field92 = 92;
double field93 = 93;
double field94 = 94;
double field95 = 95;
double field96 = 96;
double field97 = 97;
double field98 = 98;
double field99 = 99;
double field100 = 100;
}
25 changes: 25 additions & 0 deletions tests/assets/bigmessage_pb2.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit 67f8b2a

Please sign in to comment.