ByteArrayPubSubSerDe - A better default for SerDe (#101)

Author: akshaisarma

src/main/java/com/yahoo/bullet/common/BulletConfig.java

@@ -143,7 +143,7 @@ public class BulletConfig extends Config {
 
     public static final String DEFAULT_PUBSUB_CONTEXT_NAME = Context.QUERY_PROCESSING.name();
     public static final String DEFAULT_PUBSUB_CLASS_NAME = "com.yahoo.bullet.pubsub.MockPubSub";
-    public static final String DEFAULT_PUBSUB_MESSAGE_SERDE_CLASS_NAME = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
+    public static final String DEFAULT_PUBSUB_MESSAGE_SERDE_CLASS_NAME = "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe";
 
     public static final String DEFAULT_RECORD_PROVIDER_CLASS_NAME = "com.yahoo.bullet.record.avro.TypedAvroBulletRecordProvider";
 

src/main/java/com/yahoo/bullet/pubsub/ByteArrayPubSubMessageSerDe.java

@@ -0,0 +1,67 @@
+/*
+ *  Copyright 2021 Yahoo Inc.
+ *  Licensed under the terms of the Apache License, Version 2.0.
+ *  See the LICENSE file associated with the project for terms.
+ */
+package com.yahoo.bullet.pubsub;
+
+import com.yahoo.bullet.common.BulletConfig;
+import com.yahoo.bullet.common.SerializerDeserializer;
+import com.yahoo.bullet.query.Query;
+
+import java.io.Serializable;
+
+/**
+ * This SerDe is used to convert a provided {@link Query} into its serialized bytes in the {@link PubSubMessage}. When
+ * invoking the {@link PubSubMessage#getContentAsByteArray()} or {@link PubSubMessage#getContentAsQuery()}, the payload
+ * is lazily converted (and stored) as the appropriate type. This means that {@link PubSubMessage#getContent()} will
+ * return a byte[] or {@link Query} depending on what was called before it. This can be used to send a {@link Query}
+ * to the backend and if it is passed between multiple workers (i.e. serialized and deserialized multiple times), it
+ * will not needlessly convert the {@link Query} object multiple times.
+ *
+ * This behaves like the {@link IdentityPubSubMessageSerDe} for all other operations.
+ */
+public class ByteArrayPubSubMessageSerDe extends IdentityPubSubMessageSerDe {
+    private static final long serialVersionUID = -7648403271773714704L;
+
+    /**
+     * A {@link PubSubMessage} that is sticky for converting the content between byte[] and {@link Query}.
+     */
+    private static class LazyPubSubMessage extends PubSubMessage {
+        private static final long serialVersionUID = -6516915913438279870L;
+
+        private LazyPubSubMessage(String id, byte[] content, Metadata metadata) {
+            super(id, content, metadata);
+        }
+
+        @Override
+        public byte[] getContentAsByteArray() {
+            if (content instanceof Query) {
+                content = SerializerDeserializer.toBytes((Serializable) content);
+            }
+            return (byte[]) content;
+        }
+
+        @Override
+        public Query getContentAsQuery() {
+            if (content instanceof byte[]) {
+                content = SerializerDeserializer.fromBytes((byte[]) content);
+            }
+            return (Query) content;
+        }
+    }
+
+    /**
+     * Constructor.
+     *
+     * @param config The {@link BulletConfig} to configure this class.
+     */
+    public ByteArrayPubSubMessageSerDe(BulletConfig config) {
+        super(config);
+    }
+
+    @Override
+    public PubSubMessage toMessage(String id, Query query, String queryString) {
+        return toMessage(new LazyPubSubMessage(id, SerializerDeserializer.toBytes(query), new Metadata(null, queryString)));
+    }
+}

src/main/java/com/yahoo/bullet/pubsub/Metadata.java

@@ -22,11 +22,11 @@ public enum Signal {
 
     private static final long serialVersionUID = 7478596915692253699L;
     @Getter @Setter
-    private Signal signal;
+    protected Signal signal;
     // Serializable enforced through the constructor, getter, and setter. Is Object so GSON can reify an instance.
-    private Object content;
+    protected Object content;
     @Getter @Setter
-    private long created;
+    protected long created;
 
     /**
      * Default constructor that creates an empty instance of metadata.

src/main/java/com/yahoo/bullet/pubsub/PubSubMessage.java

@@ -28,11 +28,11 @@ public class PubSubMessage implements Serializable, JSONFormatter {
     public static final Charset CHARSET = StandardCharsets.UTF_8;
     private static final long serialVersionUID = 5096747716667851530L;
 
-    private String id;
+    protected String id;
     // Serializable enforced through the constructors, getter and setter. Is Object so GSON can reify an instance.
-    private Object content;
+    protected Object content;
     @Setter
-    private Metadata metadata;
+    protected Metadata metadata;
 
     /**
      * Constructor for a message having no information. Used internally. Not recommended for use.

src/main/resources/bullet_defaults.yaml

@@ -250,4 +250,4 @@ bullet.pubsub.class.name: "com.yahoo.bullet.pubsub.rest.RESTPubSub"
 # The current context. This can be QUERY_PROCESSING or QUERY_SUBMISSION. The PubSub implementation should use this to generate appropriate Publishers and Subscribers.
 bullet.pubsub.context.name: "QUERY_PROCESSING"
 # The class to use for converting and reading PubSubMessage queries sent to the backend.
-bullet.pubsub.message.serde.class.name: "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe"
+bullet.pubsub.message.serde.class.name: "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe"

src/test/java/com/yahoo/bullet/common/BulletConfigTest.java

@@ -178,7 +178,7 @@ public void testPropertiesWithPrefix() {
         BulletConfig config = new BulletConfig("src/test/resources/test_config.yaml");
         String prefix = "bullet.pubsub";
         String pubSubClassValue = "com.yahoo.bullet.pubsub.MockPubSub";
-        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
+        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe";
 
         int configSize = config.getAllWithPrefix(Optional.empty(), prefix, false).size();
         Assert.assertEquals(configSize, 5);
@@ -195,7 +195,7 @@ public void testPropertiesStripPrefix() {
         String pubsubClassKey = "class.name";
         String pubSubClassValue = "com.yahoo.bullet.pubsub.MockPubSub";
         String pubsubMessageSerDeClassKey = "message.serde.class.name";
-        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
+        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe";
 
         int configSize = config.getAllWithPrefix(Optional.empty(), prefix, true).size();
         Assert.assertEquals(configSize, 5);

src/test/java/com/yahoo/bullet/pubsub/ByteArrayPubSubMessageSerDeTest.java

@@ -0,0 +1,62 @@
+/*
+ *  Copyright 2021 Yahoo Inc.
+ *  Licensed under the terms of the Apache License, Version 2.0.
+ *  See the LICENSE file associated with the project for terms.
+ */
+package com.yahoo.bullet.pubsub;
+
+import com.yahoo.bullet.common.SerializerDeserializer;
+import com.yahoo.bullet.query.Projection;
+import com.yahoo.bullet.query.Query;
+import com.yahoo.bullet.query.Window;
+import com.yahoo.bullet.query.aggregations.AggregationType;
+import com.yahoo.bullet.query.aggregations.Raw;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class ByteArrayPubSubMessageSerDeTest {
+    @Test
+    public void testConvertingQuery() {
+        ByteArrayPubSubMessageSerDe serDe = new ByteArrayPubSubMessageSerDe(null);
+
+        Query query = new Query(new Projection(), null, new Raw(1), null, new Window(), 1L);
+        PubSubMessage actual = serDe.toMessage("id", query, "foo");
+        Assert.assertEquals(actual.getId(), "id");
+        Assert.assertEquals(actual.getContent(), SerializerDeserializer.toBytes(query));
+        Assert.assertEquals(actual.getMetadata().getContent(), "foo");
+    }
+
+    @Test
+    public void testDoNothingPubSubMessage() {
+        ByteArrayPubSubMessageSerDe serDe = new ByteArrayPubSubMessageSerDe(null);
+        PubSubMessage expected = new PubSubMessage("foo", new byte[0], new Metadata(Metadata.Signal.KILL, null));
+        Assert.assertSame(serDe.toMessage(expected), expected);
+        Assert.assertSame(serDe.fromMessage(expected), expected);
+    }
+
+    @Test
+    public void testLazyMessage() {
+        ByteArrayPubSubMessageSerDe serDe = new ByteArrayPubSubMessageSerDe(null);
+        Query query = new Query(new Projection(), null, new Raw(1), null, new Window(), 1L);
+        PubSubMessage converted = serDe.toMessage("id", query, "foo");
+        PubSubMessage reverted = serDe.fromMessage(converted);
+
+        Assert.assertSame(reverted, converted);
+
+        // Starts off as byte[]
+        Assert.assertEquals(reverted.getContent(), SerializerDeserializer.toBytes(query));
+
+        // Payload is now made a Query
+        Query revertedQuery = reverted.getContentAsQuery();
+        Assert.assertEquals(revertedQuery.getProjection().getType(), Projection.Type.PASS_THROUGH);
+        Assert.assertEquals(revertedQuery.getAggregation().getType(), AggregationType.RAW);
+        Assert.assertEquals((long) revertedQuery.getAggregation().getSize(), 1L);
+        Assert.assertEquals((long) revertedQuery.getDuration(), 1L);
+        Assert.assertSame(reverted.getContent(), revertedQuery);
+
+        // Payload is now made a byte[]
+        byte[] revertedByteArray = reverted.getContentAsByteArray();
+        Assert.assertEquals(revertedByteArray, SerializerDeserializer.toBytes(query));
+        Assert.assertSame(reverted.getContent(), revertedByteArray);
+    }
+}

src/test/java/com/yahoo/bullet/pubsub/rest/RESTPubSubConfigTest.java

@@ -123,7 +123,7 @@ public void testPropertiesWithPrefix() {
         RESTPubSubConfig config = new RESTPubSubConfig("src/test/resources/test_config.yaml");
         String prefix = "bullet.pubsub";
         String pubSubClassValue = "com.yahoo.bullet.pubsub.MockPubSub";
-        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
+        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe";
 
         int configSize = config.getAllWithPrefix(Optional.empty(), prefix, false).size();
         Assert.assertEquals(configSize, 10);
@@ -140,7 +140,7 @@ public void testPropertiesStripPrefix() {
         String pubsubClassKey = "class.name";
         String pubSubClassValue = "com.yahoo.bullet.pubsub.MockPubSub";
         String pubsubMessageSerDeClassKey = "message.serde.class.name";
-        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
+        String pubSubMessageSerDeClassValue = "com.yahoo.bullet.pubsub.ByteArrayPubSubMessageSerDe";
 
         int configSize = config.getAllWithPrefix(Optional.empty(), prefix, true).size();
         Assert.assertEquals(configSize, 10);