[dvc][server][controller][samza] Heap size estimation improvement #1281

FelixGV · 2024-11-04T22:25:00Z

[dvc][server][controller][samza] Heap size estimation improvement

Introduced two new utilities to make our on-heap memory usage assessment more
accurate, and easier to maintain as class hierarchies evolve:

ClassSizeEstimator: Predicts (based on assumptions about how the JVM's
memory layout works) the shallow size of a class. This includes the object
header, all primitive fields, and all references to other objects (but it
does not count these other objects, hence the shallowness). Reflection is
used in this class.
InstanceSizeEstimator: Predicts the size of instances of a limited number
of classes. This is not a general-purpose utility, and it requires some
manual effort to onboard a new class. Reflection is not used in this class.

The general design goals are the following:

Reflection should only be used once per class per runtime, and the result
of this logic should be stored in static constants.
On the hot path, there should be no reflection, and we should leverage our
knowledge of the Venice code base to determine which objects are meant to
be counted or not. For example, singleton or otherwise shared instances
should not be counted, since their amortized cost is negligible (besides
the size of the pointer to refer to them).

The above utilities have been integrated in all classes that implement the
Measurable interface, and several new classes have been given this interface
as well. The Measurable::getSize function has been renamed getHeapSize, to
minimize the chance that it could clash with other function names, and to
make it extra clear what kind of size is meant.

Miscellaneous:

Fixed NPEs in AdminConsumptionTask::executeMessagesAndCollectResults.
Minor efficiency improvements to PubSubMessageHeaders and ApacheKafkaUtils
so that empty headers (a common case) carry less overhead. Also made the
PubSubMessageHeaders implement Iterable.
Created a DefaultLeaderMetadata static class in VeniceWriter, so that a
shared instance can be leveraged in cases where that object is always the
same (e.g. when producing to the RT topic).
BlobSnapshotManagerTest improvements:
- Added timeouts to all tests.
- Fixed a race condition in testMultipleThreads.

How was this PR tested?

New unit tests.

Does this PR introduce any user-facing changes?

No. You can skip the rest of this section.
Yes. Make sure to explain your proposed changes and call out the behavior change.

gaojieliu · 2024-11-05T17:56:07Z

Have you got a chance to take a look at openjdk/jol implementation?
https://github.com/openjdk/jol
It might carry some issues, but I guess we can analyze the strategy this lib is using to see whether we can borrow some or not.

FelixGV · 2024-11-05T20:53:42Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java

+/**
+ * Utility class to help in implementing {@link Measurable#getSize()}. A couple of important points:
+ *
+ * 1. This utility class does not "measure" the heap size, but rather attempts to "predict" it, based on knowledge of
+ *    the internals of the JVM. If any of the assumptions are wrong, then of course the results will be inaccurate.
+ * 2. This utility class assumes we are using the HotSpot JVM.
+ */
+public class HeapSizeEstimator {


Starting a thread to respond to @gaojieliu's general comment:

Have you got a chance to take a look at openjdk/jol implementation?
https://github.com/openjdk/jol
It might carry some issues, but I guess we can analyze the strategy this lib is using to see whether we can borrow some or not.

I posted this work on social media and received a lot of interesting examples from folks about other open source projects with the same or similar purpose. I'll maintain a list of these here:

https://github.com/openjdk/jol

https://github.com/ehcache/sizeof

https://github.com/apache/lucene/blob/main/lucene/core/src/java/org/apache/lucene/util/RamUsageEstimator.java

https://github.com/jbellis/jamm

For now, this work looks pretty similar to (at least some parts of) the above, but I think there is still potential to have a good implementation as part of Venice. In particular, I want a solution in which there is no hot path reflection. For now, this tradeoff is not apparent because I have not yet included examples of how we can leverage this utility in the rest of our code, but I will add it soon when the basic functionality is stabilized...

The reflection mentioned here is mainly about performance, right?
Have you tried these alternatives and run it with JMH?
I tried jol and the initialization is slow, but after that, it is fairly fast.

I did a small test against jol and it seems it can work properly with Oracle OpenJDK8 release, but not MSFT OpenJDK11 release.

@alex-dubrouski
Do you know why MSFT-jdk11 won't work well with attached agent to measure the object size?
I tried to add this jvm arg, but it still doesn't work:

-Djdk.attach.allowAttachSelf=true

And both jol and sizeof are suffering from the same issue with JDK11 as they are using the similar technologies..

I will check a bit later

Well, Felix, that's why we should probably use JOL :D ;)

But where is the learning value in that, @alex-dubrouski :D ?

Besides, this memory estimation code needs to run in DVC, which is one of our client libs, meaning we cannot tightly control which JVM it's going to run in. Ideally, it should not depend on agents or finicky Unsafe code...

That being said, using JOL in unit test to validate the accuracy of this utility might be a good idea 🤔

BTW... you did not answer my question @alex-dubrouski 😅 but I found out that the hashIsZero was added 5 years ago: openjdk/jdk@89a267c

According to JDK-8221836, this change got released in Java 13, which is weird because I thought I was looking at the Java 17 source code in my IDE, but apparently not; it was probably hanging on to the source of an older Java version.

And indeed, I can get it to show me the proper source and see that it's there now.

I tested on JDK17 so yeah, it is there. Plus it is a field not method.

alex-dubrouski · 2024-11-06T17:15:44Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java

+
+  /** Deal with alignment by rounding up to the nearest alignment boundary. */
+  private static int roundUpToNearestAlignment(int size) {
+    int partialAlignmentWindowUsage = size % ALIGNMENT_SIZE;


Alignment could be modified with command line.

I know that it can be configured at start time. But are you saying it can be modified at run time?

No, it can't be changed at runtime, but you defined static alignment which might not be accurate. That's an edge case of course, I was just saying.

Oh I see... you mean that alignment can be overridden by a JVM option, in which case it wouldn't decide just based on whether it's a 32 bits or 64 bits VM?

Is it via -XX:ObjectAlignmentInBytes?

Yes, it is set statically at bootstrap, but code does not take this into account. Feel free to ignore that, since this is very rare edge case.

alex-dubrouski · 2024-11-06T17:16:42Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java

+     * {@link Class#getDeclaredFields()} returns the fields (of all visibility, from private to public and everything
+     * in between) of the class, but not of its parent class.
+     */
+    for (Field f: c.getDeclaredFields()) {


Some of the fields are not visible. For example there are native references.

Could you point me to a class (perhaps from the JDK) that has a native reference in it? I'd like to test that scenario and see what happens...

Java.lang.Object :)

Then I'm definitely not taking it into account 😅 ... where can I learn more about this?

Alex Shipilev explained that in one of the videos (video is not in English...)

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java

gaojieliu · 2024-11-07T17:52:00Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java

+   * @throws {@link StackOverflowError} It should be noted that this function is recursive in nature and so any class
+   *         which contains a loop in its class graph will cause a stack overflow.
+   */
+  public static <T> int getClassOverhead(Class<T> c) {


So this function is mainly calculating the class overhead, not the deep size, right?
If the object contains a nested hashmap, this class will only measure the overhead of the class, and the user needs to manually add the class overhead on top of the hashmap cost, which needs to be calculated manually?

I think it will be useful to have a function to measure both class overhead and the actual data.

The above requirement is mainly about measuring the deep size of the object.

If we do want to implement such feature, it will be great if we can have some filtering logic to ignore certain types of fields, which can be referring to some common references.

Yes, this function does not measure the deep size of an object, and that is also the reason why it takes a Class param and not an Object. Essentially, my idea is to hand-code the "variable" part of the measurement of the classes of interest, and the function in this utility is just to make it easier to have the "base cost" (the non-variable part). This utility is not intended to be called in the hot path.

I will integrate the usage of this utility into the MemoryBoundBlockingQueue measurement code. I have not done so yet because I want to be extra sure the utility is as accurate as possible (which it may or may not be, as I am not sure if I have chased all of the bugs properly yet, but at least it should be pretty close...).

Introduced a new HeapSizeEstimator utility to make our on-heap memory usage assessment more accurate, and easier to maintain as class hierarchies evolve. For now it is only used in a unit test, not yet in the main code.

… measures the allocated memory.

Also did some refactorings and added tests.

HeapSizeEstimator renamed to ClassSizeEstimator MeasurableUtils renamed to InstanceSizeEstimator HeapSizeEstimatorTest is now abstract, with two subclasses: ClassSizeEstimatorTest and InstanceSizeEstimatorTest.

Simplified the class size measurement by removing the option of recursing into the class graph. It turns out that after integrating into all parts of the main code that need heap size measurement, this was never used, so the utility will now provide shallow measurement exclusively. Miscellaneous: - Minor efficiency improvements to PubSubMessageHeaders, especially when they are empty. Also made them implement Iterable, which is more efficient than creating a temporary List to iterate over. Added unit tests for this.

…sults

… fixed a race condition in testMultipleThreads.

...roller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java

Most of these changes are just IDE hints about potential NPEs, and fields which are unused, can be made final or made into a local variable.

gaojieliu

Looks good overall, left some minor comments.

gaojieliu · 2024-11-13T01:11:27Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/ClassSizeEstimator.java

+  private static int roundUpToNearest(int size, int intervalSize) {
+    int partialAlignmentWindowUsage = size % intervalSize;
+    int waste = partialAlignmentWindowUsage == 0 ? 0 : intervalSize - partialAlignmentWindowUsage;
+    int finalSize = size + waste;
+    return finalSize;


Is this function similar to the following?

Math.ceil((double)size / intervalSize) * intervalSize

That's right, and that is the logic I'm using in HeapSizeEstimatorTest::roundUpToNearestAlignment, so they are fully equivalent AFAICT. We could pick either one for the main code. I thought this one was slightly more clear, so I picked it for the main code, but we could do the reverse. LMK if you have a preference.

Using the ceil function is common for roundup, I guess it is good to use the same logic in both main and test code.
This is a very minor concern as the logic is simple.
You can decide on your own.

gaojieliu · 2024-11-13T01:20:36Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/InstanceSizeEstimator.java

+        return value -> getSize((KafkaMessageEnvelope) value);
+      } else if (ProducerMetadata.class.isAssignableFrom(type)) {
+        return value -> getSize((ProducerMetadata) value);
+      } else if (ByteBuffer.class.isAssignableFrom(type)) {


If some class extends ByteBuffer to add some additional fields, the returned function won't count them, right?

That's right. For now, that function is only meant to be used with HeapByteBuffer instances and will throw IllegalArgumentException if a DirectBB is passed. If a subclass of HeapBB is passed, there could be imprecision. I'm not aware of such case though. Are you?

So can we just use HeapByteBuffer here?

gaojieliu · 2024-11-13T17:29:38Z

internal/venice-common/src/main/java/com/linkedin/venice/memory/ClassSizeEstimator.java

+
+    // Iterate from the end to the beginning, so we go from parent to sub
+    for (int i = classHierarchyFromSubclassToParent.size() - 1; i >= 0; i--) {
+      int classFieldsOverhead = overheadOfFields(classHierarchyFromSubclassToParent.get(i));


Do we want to use the cache here? Some parent class might be calculated before.

I think that would be incorrect, because the cache contains the final size including object header and alignment. Here, we are looking for the "intermediate result" (the fields only), so we can add all these intermediate results together and account for alignment afterwards. We could have a separate cache for these intermediate results but I don't think it's worth it given that the intent is to only call this once per class per runtime, never on the hot path...

Got it, makes sense.
Can you also convert these explanation into Javadoc?

internal/venice-common/src/main/java/com/linkedin/venice/memory/ClassSizeEstimator.java

gaojieliu · 2024-11-13T17:36:03Z

...ts/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreBufferService.java

  }

  private static class LeaderQueueNode extends QueueNode {
+    private static final int SHALLOW_CLASS_OVERHEAD = ClassSizeEstimator.getClassOverhead(QueueNode.class);


Why not LeaderQueueNode?

Good catch! And that is one of the cases I didn't write a unit test for... I think I need to write one now to have a clean conscience 😅

stuck forever if some task for it got scheduled but then got cancelled before it started executing...

FelixGV commented Nov 5, 2024

View reviewed changes

FelixGV force-pushed the memory_usage_estimation branch from 94ff779 to 626a47d Compare November 5, 2024 20:54

alex-dubrouski reviewed Nov 6, 2024

View reviewed changes

internal/venice-common/src/main/java/com/linkedin/venice/memory/HeapSizeEstimator.java Outdated Show resolved Hide resolved

FelixGV force-pushed the memory_usage_estimation branch 2 times, most recently from 16ae593 to c39e456 Compare November 7, 2024 16:49

gaojieliu reviewed Nov 7, 2024

View reviewed changes

FelixGV force-pushed the memory_usage_estimation branch from b635dff to 7b9581a Compare November 12, 2024 17:12

FelixGV changed the title ~~[test] HeapSizeEstimator~~ [dvc][server][controller][samza] Heap size estimation improvement Nov 12, 2024

FelixGV added 16 commits November 12, 2024 13:19

[test] HeapSizeEstimator

f3ad501

Introduced a new HeapSizeEstimator utility to make our on-heap memory usage assessment more accurate, and easier to maintain as class hierarchies evolve. For now it is only used in a unit test, not yet in the main code.

Fixed many bugs discovered as part of a new test which experimentally…

1dda2ba

… measures the allocated memory.

Fixed bug related to "superclass gaps", and improved testing.

1cc7688

Fixed static analysis and made measurement a little more lenient.

da3e3b8

Added memoization, as suggest by Alex D., and refactored a little bit.

c748f17

Static analysis appeasement.

5a10245

More tests and minor cleanups...

e6d2b6d

Integrated new measurement code throughout the main code.

2af2a8c

Also did some refactorings and added tests.

Minor fixes to unit tests and static analysis.

3c8a5aa

Many cleanups

a1d68c7

HeapSizeEstimator renamed to ClassSizeEstimator MeasurableUtils renamed to InstanceSizeEstimator HeapSizeEstimatorTest is now abstract, with two subclasses: ClassSizeEstimatorTest and InstanceSizeEstimatorTest.

More static analysis appeasement.

558653d

Fixed NPE in AdminConsumptionTask::executeMessagesAndCollectResults

d00365b

One more NPE fix in AdminConsumptionTask::executeMessagesAndCollectRe…

93b3dce

…sults

BlobSnapshotManagerTest improvements: added timeouts to all tests and…

dee6b6f

… fixed a race condition in testMultipleThreads.

Static analysis appeasement.

d5573b2

FelixGV force-pushed the memory_usage_estimation branch from 70a7ee0 to d5573b2 Compare November 12, 2024 21:21

xunyin8 reviewed Nov 12, 2024

View reviewed changes

...roller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java Outdated Show resolved Hide resolved

...roller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java Outdated Show resolved Hide resolved

More cleanups to the AdminConsumptionTask and related classes.

2cc9111

Most of these changes are just IDE hints about potential NPEs, and fields which are unused, can be made final or made into a local variable.

gaojieliu reviewed Nov 13, 2024

View reviewed changes

Fixed one more edge case in AdminConsumptionTask where a store could get

605a415

stuck forever if some task for it got scheduled but then got cancelled before it started executing...

[dvc][server][controller][samza] Heap size estimation improvement #1281

Are you sure you want to change the base?

[dvc][server][controller][samza] Heap size estimation improvement #1281

Conversation

FelixGV commented Nov 4, 2024 • edited Loading

How was this PR tested?

Does this PR introduce any user-facing changes?

gaojieliu commented Nov 5, 2024

FelixGV Nov 5, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gaojieliu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

FelixGV commented Nov 4, 2024 •

edited

Loading

FelixGV Nov 5, 2024 •

edited

Loading