[Bug]: AvroCoder uses a lot of memory because it keeps instantiating the same DatumReader and DatumWriter

[wollowizard]

What happened?

In a dataflow streaming pipeline I noticed that around 28% of the memory allocations are made by AvroDatumFactory$ReflectDatumFactory.apply which keeps being called with the same schemas as input (see figure)

This is part of the AvroCoder.decode call, which in theory has a mechanism to avoid exactly this situation: it has a ThreadLocal to cache the previously computed BinaryDecoder.

However, it seems that this ThreadLocal is always empty and so the value needs to be recomputed each time. It is unclear to me at the moment if this ThreadLocal is empty because the runner keeps creating new threads, or because the ThreadLocal, which is actually an EmptyOnDeserializationThreadLocal, is always empty when deserialization happens.

Recreating the BinaryDecoder's initial value means that the AvroDatumFactory is called again and again, with the same Schemas as inputs. This seems wasteful, even more considering that the DatumReader and DatumWriter should be reusable and thread-safe.

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[liferoad]

cc @scwhittle

[wollowizard]

.take-issue

[scwhittle]

The heaps stacks are a snapshot of memory used and where it was allocated, not the rate of allocations. So I don't think it show that the allocations are happening repeatedly for a single thread, just that the memory overhead of allocating per-thread adds up (since there can be 300-500 threads).

If the DatumReader and DatumWriter are thread-safe, can we just remove the thread-locals used for them here? https://github.com/apache/beam/blob/master/sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java#L380

There still may be some benefit of a static cache somewhere (since there may be multiple AvroCoder instances that are the same) but I think we can see if just removing the thread-local helps.

[wollowizard]

hi @scwhittle that was my first try but I seem to remember that it didn't work because the reader and writer are not serializable. I guess that is also the reason why an EmptyOnDeserializationThreadLocalwas used, instead of a plain ThreadLocal.
Would you like me to double check that?

Also, yes I confirm there are places even in the beam sdk where the same AvroCoder is instantiated multiple times instead of being cached, for example here (those multiple allocations are visible in the memory allocation flame graph too)

[scwhittle]

I was thinking that most coder usage was setup during graph construction and thus new coders weren't being allocated, but it does appear there are AvroCoder.of sprinkled through various sites. So your original hypothesis that these reader/writers were being created regularly seems true, since the thread-local cache is per-AvroCoder instance and we are creating new ones all the time. Your example checkpointcoder one could be changed to use a singleton like some other UnboundedSources but fixing generically seems worthwhile to avoid this pain point.

The proposed approach will fix that we have duplicate instances of the reader/writer but we will still end up having multiple AvroCoder instances which are logically the same and thus can have multiple copies per-thread of these encoder/decoder objects (and we won't resuse existing ones for a freshly created coder).
https://source.corp.google.com/piper///depot/google3/third_party/java_src/apache_beam4g/srcs/sdks/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/coders/AvroCoder.java;l=376;bpv=1;bpt=1;rcl=755306863

So it seems it could be better to move the static cache to within AvroCoder itself so that all of the state for the same class/datumfactory is reused.

To clarify, I was suggesting removing the threadlocal for reader/writer in AvroCoder if they are thread-safe (which it seems they must be if we are changing to reuse them across threads with the cache). There is no reason to have a per-thread writer that is actually the same object for all the threads due to the cache. So I think it could just be transient DatumReader that we either restore in deserialization (example https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/util/WindowedValue.java#L892) or initialize lazily (example https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/SerializableCoder.java#L178)

[wollowizard]

If I understand well, you propose to have a static cache within AvroCoder to cache a full AvroCoder, right?

I think this indeed solves the problem of people creating new AvroCoders outside of graph setup. And this not just in the sdk but also in user code (actually we had this exact issue in our own code base too, that's where my investigation starts from).

For example a developer writing a coder for kafka events can very easily end up creating a new AvroCoder at every received kafka message and the costs associated with it are unclear to a developer that hasn't run a profiler with a realistic workload.

@scwhittle Would you like me to open a PR to do that? Also, would you consider merging this PR first so it has a chance of being included in the next beam release, and in parallel open another issue to work on your proposal?

[scwhittle]

If I understand well, you propose to have a static cache within AvroCoder to cache a full AvroCoder, right?

Yup! as well as getting rid of those unnecessary thread-locals for the reader/writer as even if we have the AvroCoder cache that will duplicate those per-thread.

Would you like me to open a PR to do that? Also, would you consider merging this PR first so it has a chance of being included in the next beam release, and in parallel open another issue to work on your proposal?

I believe that we have about a month before the next beam release cut. How about keeping the existing PR and trying the new approach in a separate PR? I don't think it will be much more complicated than the current PR so hopefully won't take long to code up. If we start getting close to deadline we can just merge the existing PR.

[wollowizard]

Agree on everything 👍 I'll send another PR then