WFLY-18066 Add ability to calculate buffer size of ByteBufferMarshall…

…edValue without generating its buffer.

clustering/marshalling/protostream/src/main/java/org/wildfly/clustering/marshalling/protostream/ByteBufferMarshalledKeyMarshaller.java

@@ -24,6 +24,7 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.OptionalInt;
 
 import org.infinispan.protostream.descriptors.WireType;
 import org.wildfly.clustering.marshalling.spi.ByteBufferMarshalledKey;
@@ -70,6 +71,11 @@ public void writeTo(ProtoStreamWriter writer, ByteBufferMarshalledKey<Object> ke
         }
     }
 
+    @Override
+    public OptionalInt size(ProtoStreamOperation context, ByteBufferMarshalledKey<Object> value) {
+        return value.size();
+    }
+
     @SuppressWarnings("unchecked")
     @Override
     public Class<? extends ByteBufferMarshalledKey<Object>> getJavaClass() {

clustering/marshalling/protostream/src/main/java/org/wildfly/clustering/marshalling/protostream/ByteBufferMarshalledValueMarshaller.java

@@ -0,0 +1,75 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2023, Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags. See the copyright.txt file in the
+ * distribution for a full listing of individual contributors.
+ *
+ * This is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation; either version 2.1 of
+ * the License, or (at your option) any later version.
+ *
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this software; if not, write to the Free
+ * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
+ * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
+ */
+
+package org.wildfly.clustering.marshalling.protostream;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.OptionalInt;
+
+import org.infinispan.protostream.descriptors.WireType;
+import org.wildfly.clustering.marshalling.spi.ByteBufferMarshalledValue;
+
+/**
+ * {@link ProtoStreamMarshaller} for a {@link ByteBufferMarshalledValue}.
+ * @author Paul Ferraro
+ */
+public class ByteBufferMarshalledValueMarshaller implements ProtoStreamMarshaller<ByteBufferMarshalledValue<Object>> {
+
+    private static final int BUFFER_INDEX = 1;
+
+    @Override
+    public ByteBufferMarshalledValue<Object> readFrom(ProtoStreamReader reader) throws IOException {
+        ByteBuffer buffer = null;
+        while (!reader.isAtEnd()) {
+            int tag = reader.readTag();
+            switch (WireType.getTagFieldNumber(tag)) {
+                case BUFFER_INDEX:
+                    buffer = reader.readByteBuffer();
+                    break;
+                default:
+                    reader.skipField(tag);
+            }
+        }
+        return new ByteBufferMarshalledValue<>(buffer);
+    }
+
+    @Override
+    public void writeTo(ProtoStreamWriter writer, ByteBufferMarshalledValue<Object> key) throws IOException {
+        ByteBuffer buffer = key.getBuffer();
+        if (buffer != null) {
+            writer.writeBytes(BUFFER_INDEX, buffer.mark());
+            buffer.reset();
+        }
+    }
+
+    @Override
+    public OptionalInt size(ProtoStreamOperation context, ByteBufferMarshalledValue<Object> value) {
+        return value.size();
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public Class<? extends ByteBufferMarshalledValue<Object>> getJavaClass() {
+        return (Class<ByteBufferMarshalledValue<Object>>) (Class<?>) ByteBufferMarshalledValue.class;
+    }
+}

clustering/marshalling/protostream/src/main/java/org/wildfly/clustering/marshalling/protostream/MarshallingMarshallerProvider.java

@@ -22,16 +22,12 @@
 
 package org.wildfly.clustering.marshalling.protostream;
 
-import java.nio.ByteBuffer;
-
-import org.wildfly.clustering.marshalling.spi.ByteBufferMarshalledValue;
-
 /**
  * @author Paul Ferraro
  */
 public enum MarshallingMarshallerProvider implements ProtoStreamMarshallerProvider {
     BYTE_BUFFER_MARSHALLED_KEY(new ByteBufferMarshalledKeyMarshaller()),
-    BYTE_BUFFER_MARSHALLED_VALUE(new FunctionalScalarMarshaller<>(Scalar.BYTE_BUFFER.cast(ByteBuffer.class), () -> new ByteBufferMarshalledValue<>(null), ByteBufferMarshalledValue::isEmpty, ByteBufferMarshalledValue::getBuffer, ByteBufferMarshalledValue::new)),
+    BYTE_BUFFER_MARSHALLED_VALUE(new ByteBufferMarshalledValueMarshaller()),
     ;
     private final ProtoStreamMarshaller<?> marshaller;
 

clustering/marshalling/spi/src/main/java/org/wildfly/clustering/marshalling/spi/ByteBufferMarshalledValue.java

@@ -28,6 +28,7 @@
 import java.io.Serializable;
 import java.nio.ByteBuffer;
 import java.util.Objects;
+import java.util.OptionalInt;
 
 import org.jboss.logging.Logger;
 
@@ -64,11 +65,18 @@ public synchronized boolean isEmpty() {
 
     public synchronized ByteBuffer getBuffer() throws IOException {
         ByteBuffer buffer = this.buffer;
-        if (buffer != null) return buffer;
-        if (this.object == null) return null;
-        ByteBuffer result = this.marshaller.write(this.object);
-        LOGGER.debugf("Marshalled size of %s(%s) = %d bytes", this.object.getClass().getCanonicalName(), this.object, result.limit() - result.arrayOffset());
-        return result;
+        if ((buffer == null) && (this.object != null)) {
+            // Since the wrapped object is likely mutable, we cannot cache the generated buffer
+            buffer = this.marshaller.write(this.object);
+            // N.B. Refrain from logging wrapped object
+            // If wrapped object contains an EJB proxy, toString() will trigger an EJB invocation!
+            LOGGER.debugf("Marshalled size of %s object = %d bytes", this.object.getClass().getCanonicalName(), buffer.limit() - buffer.arrayOffset());
+        }
+        return buffer;
+    }
+
+    public synchronized OptionalInt size() {
+        return (this.buffer != null) ? OptionalInt.of(this.buffer.limit() - this.buffer.arrayOffset()) : this.marshaller.size(this.object);
     }
 
     @SuppressWarnings("unchecked")
@@ -110,10 +118,9 @@ public boolean equals(Object object) {
 
     @Override
     public String toString() {
-        Object object = this.object;
-        if (object != null) return object.toString();
-        ByteBuffer buffer = this.buffer;
-        return (buffer != null) ? buffer.toString() : null;
+        // N.B. Refrain from logging wrapped object
+        // If wrapped object contains an EJB proxy, toString() will trigger an EJB invocation!
+        return String.format("%s [%s]", this.getClass().getName(), (this.object != null) ? this.object.getClass().getName() : "<serialized>");
     }
 
     private void writeObject(ObjectOutputStream out) throws IOException {