helix-core/src/main/java/org/apache/helix/manager/zk/ZNRecordStreamingSerializer.java - helix - Git at Google

 package org.apache.helix.manager.zk;

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;

 import org.I0Itec.zkclient.exception.ZkMarshallingError;
 import org.I0Itec.zkclient.serialize.ZkSerializer;
 import org.apache.commons.codec.binary.Base64;
 import org.apache.helix.HelixException;
 import org.apache.helix.ZNRecord;
 import org.apache.helix.util.GZipCompressionUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.codehaus.jackson.JsonFactory;
 import org.codehaus.jackson.JsonGenerator;
 import org.codehaus.jackson.JsonParser;
 import org.codehaus.jackson.JsonToken;

 import com.google.common.collect.Maps;

 public class ZNRecordStreamingSerializer implements ZkSerializer {
   private static Logger LOG = LoggerFactory.getLogger(ZNRecordStreamingSerializer.class);

   private static int getListFieldBound(ZNRecord record) {
     int max = Integer.MAX_VALUE;
     if (record.getSimpleFields().containsKey(ZNRecord.LIST_FIELD_BOUND)) {
       String maxStr = record.getSimpleField(ZNRecord.LIST_FIELD_BOUND);
       try {
         max = Integer.parseInt(maxStr);
       } catch (Exception e) {
         LOG.error("IllegalNumberFormat for list field bound: " + maxStr);
       }
     }
     return max;
   }

   @Override
   public byte[] serialize(Object data) throws ZkMarshallingError {
     if (!(data instanceof ZNRecord)) {
       // null is NOT an instance of any class
       LOG.error("Input object must be of type ZNRecord but it is " + data
           + ". Will not write to zk");
       throw new HelixException("Input object is not of type ZNRecord (was " + data + ")");
     }

     // apply retention policy on list field
     ZNRecord record = (ZNRecord) data;
     int max = getListFieldBound(record);
     if (max < Integer.MAX_VALUE) {
       Map<String, List<String>> listMap = record.getListFields();
       for (String key : listMap.keySet()) {
         List<String> list = listMap.get(key);
         if (list.size() > max) {
           listMap.put(key, list.subList(0, max));
         }
       }
     }
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     byte[] serializedBytes = null;
     try {
       JsonFactory f = new JsonFactory();
       JsonGenerator g = f.createJsonGenerator(baos);

       g.writeStartObject();

       // write id field
       g.writeRaw("\n  ");
       g.writeStringField("id", record.getId());

       // write simepleFields
       g.writeRaw("\n  ");
       g.writeObjectFieldStart("simpleFields");
       for (String key : record.getSimpleFields().keySet()) {
         g.writeRaw("\n    ");
         g.writeStringField(key, record.getSimpleField(key));
       }
       g.writeRaw("\n  ");
       g.writeEndObject(); // for simpleFields

       // write listFields
       g.writeRaw("\n  ");
       g.writeObjectFieldStart("listFields");
       for (String key : record.getListFields().keySet()) {
         // g.writeStringField(key, record.getListField(key).toString());

         // g.writeObjectFieldStart(key);
         g.writeRaw("\n    ");
         g.writeArrayFieldStart(key);
         List<String> list = record.getListField(key);
         for (String listValue : list) {
           g.writeString(listValue);
         }
         // g.writeEndObject();
         g.writeEndArray();

       }
       g.writeRaw("\n  ");
       g.writeEndObject(); // for listFields

       // write mapFields
       g.writeRaw("\n  ");
       g.writeObjectFieldStart("mapFields");
       for (String key : record.getMapFields().keySet()) {
         // g.writeStringField(key, record.getMapField(key).toString());
         g.writeRaw("\n    ");
         g.writeObjectFieldStart(key);
         Map<String, String> map = record.getMapField(key);
         for (String mapKey : map.keySet()) {
           g.writeRaw("\n      ");
           g.writeStringField(mapKey, map.get(mapKey));
         }
         g.writeRaw("\n    ");
         g.writeEndObject();

       }
       g.writeRaw("\n  ");
       g.writeEndObject(); // for mapFields

       byte[] rawPayload = record.getRawPayload();
       if (rawPayload != null && rawPayload.length > 0) {
         // write rawPayload
         g.writeRaw("\n  ");
         g.writeStringField("rawPayload", new String(Base64.encodeBase64(rawPayload), "UTF-8"));
       }

       g.writeRaw("\n");
       g.writeEndObject(); // for whole znrecord

       // important: will force flushing of output, close underlying output
       // stream
       g.close();
       serializedBytes = baos.toByteArray();
       // apply compression if needed
       if (record.getBooleanField("enableCompression", false) || serializedBytes.length > ZNRecord.SIZE_LIMIT) {
         serializedBytes = GZipCompressionUtil.compress(serializedBytes);
       }
     } catch (Exception e) {
       LOG.error("Exception during data serialization. Will not write to zk. Data (first 1k): "
           + new String(baos.toByteArray()).substring(0, 1024), e);
       throw new HelixException(e);
     }
     // check size
     if (serializedBytes.length > ZNRecord.SIZE_LIMIT) {
       LOG.error("Data size larger than 1M, ZNRecord.id: " + record.getId()
           + ". Will not write to zk. Data (first 1k): "
           + new String(serializedBytes).substring(0, 1024));
       throw new HelixException("Data size larger than 1M, ZNRecord.id: " + record.getId());
     }

     return serializedBytes;
   }

   @Override
   public Object deserialize(byte[] bytes) throws ZkMarshallingError {
     if (bytes == null || bytes.length == 0) {
       LOG.error("ZNode is empty.");
       return null;
     }

     ByteArrayInputStream bais = new ByteArrayInputStream(bytes);

     ZNRecord record = null;
     String id = null;
     Map<String, String> simpleFields = Maps.newHashMap();
     Map<String, List<String>> listFields = Maps.newHashMap();
     Map<String, Map<String, String>> mapFields = Maps.newHashMap();
     byte[] rawPayload = null;

     try {
       // decompress the data if its already compressed
       if (GZipCompressionUtil.isCompressed(bytes)) {
         byte[] uncompressedBytes = GZipCompressionUtil.uncompress(bais);
         bais = new ByteArrayInputStream(uncompressedBytes);
       }
       JsonFactory f = new JsonFactory();
       JsonParser jp = f.createJsonParser(bais);

       jp.nextToken(); // will return JsonToken.START_OBJECT (verify?)
       while (jp.nextToken() != JsonToken.END_OBJECT) {
         String fieldname = jp.getCurrentName();
         jp.nextToken(); // move to value, or START_OBJECT/START_ARRAY
         if ("id".equals(fieldname)) {
           // contains an object
           id = jp.getText();
         } else if ("simpleFields".equals(fieldname)) {
           while (jp.nextToken() != JsonToken.END_OBJECT) {
             String key = jp.getCurrentName();
             jp.nextToken(); // move to value
             simpleFields.put(key, jp.getText());
           }
         } else if ("mapFields".equals(fieldname)) {
           // user.setVerified(jp.getCurrentToken() == JsonToken.VALUE_TRUE);
           while (jp.nextToken() != JsonToken.END_OBJECT) {
             String key = jp.getCurrentName();
             mapFields.put(key, new TreeMap<String, String>());
             jp.nextToken(); // move to value

             while (jp.nextToken() != JsonToken.END_OBJECT) {
               String mapKey = jp.getCurrentName();
               jp.nextToken(); // move to value
               mapFields.get(key).put(mapKey, jp.getText());
             }
           }

         } else if ("listFields".equals(fieldname)) {
           // user.setUserImage(jp.getBinaryValue());
           while (jp.nextToken() != JsonToken.END_OBJECT) {
             String key = jp.getCurrentName();
             listFields.put(key, new ArrayList<String>());
             jp.nextToken(); // move to value
             while (jp.nextToken() != JsonToken.END_ARRAY) {
               listFields.get(key).add(jp.getText());
             }

           }

         } else if ("rawPayload".equals(fieldname)) {
           rawPayload = Base64.decodeBase64(jp.getText());
         } else {
           throw new IllegalStateException("Unrecognized field '" + fieldname + "'!");
         }
       }
       jp.close(); // ensure resources get cleaned up timely and properly

       if (id == null) {
         throw new IllegalStateException("ZNRecord id field is required!");
       }
       record = new ZNRecord(id);
       record.setSimpleFields(simpleFields);
       record.setListFields(listFields);
       record.setMapFields(mapFields);
       record.setRawPayload(rawPayload);
     } catch (Exception e) {
       LOG.error("Exception during deserialization of bytes: " + new String(bytes), e);
     }
     return record;
   }

   public static void main(String[] args) {
     ZNRecord record = new ZNRecord("record");
     final int recordSize = 10;
     for (int i = 0; i < recordSize; i++) {
       record.setSimpleField("" + i, "" + i);
       record.setListField("" + i, new ArrayList<String>());
       for (int j = 0; j < recordSize; j++) {
         record.getListField("" + i).add("" + j);
       }

       record.setMapField("" + i, new TreeMap<String, String>());
       for (int j = 0; j < recordSize; j++) {
         record.getMapField("" + i).put("" + j, "" + j);
       }
     }

     ZNRecordStreamingSerializer serializer = new ZNRecordStreamingSerializer();
     byte[] bytes = serializer.serialize(record);
     System.out.println(new String(bytes));
     ZNRecord record2 = (ZNRecord) serializer.deserialize(bytes);
     System.out.println(record2);

     long start = System.currentTimeMillis();
     for (int i = 0; i < 100; i++) {
       bytes = serializer.serialize(record);
       // System.out.println(new String(bytes));
       record2 = (ZNRecord) serializer.deserialize(bytes);
       // System.out.println(record2);
     }
     long end = System.currentTimeMillis();
     System.out.println("ZNRecordStreamingSerializer time used: " + (end - start));

     ZNRecordSerializer serializer2 = new ZNRecordSerializer();
     bytes = serializer2.serialize(record);
     // System.out.println(new String(bytes));
     record2 = (ZNRecord) serializer2.deserialize(bytes);
     // System.out.println(record2);

     start = System.currentTimeMillis();
     for (int i = 0; i < 100; i++) {
       bytes = serializer2.serialize(record);
       // System.out.println(new String(bytes));
       record2 = (ZNRecord) serializer2.deserialize(bytes);
       // System.out.println(record2);
     }
     end = System.currentTimeMillis();
     System.out.println("ZNRecordSerializer time used: " + (end - start));

   }
 }
	package org.apache.helix.manager.zk;

	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	import java.io.ByteArrayInputStream;
	import java.io.ByteArrayOutputStream;
	import java.io.StringWriter;
	import java.util.ArrayList;
	import java.util.List;
	import java.util.Map;
	import java.util.TreeMap;
	import java.util.zip.GZIPInputStream;
	import java.util.zip.GZIPOutputStream;

	import org.I0Itec.zkclient.exception.ZkMarshallingError;
	import org.I0Itec.zkclient.serialize.ZkSerializer;
	import org.apache.commons.codec.binary.Base64;
	import org.apache.helix.HelixException;
	import org.apache.helix.ZNRecord;
	import org.apache.helix.util.GZipCompressionUtil;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;
	import org.codehaus.jackson.JsonFactory;
	import org.codehaus.jackson.JsonGenerator;
	import org.codehaus.jackson.JsonParser;
	import org.codehaus.jackson.JsonToken;

	import com.google.common.collect.Maps;

	public class ZNRecordStreamingSerializer implements ZkSerializer {
	private static Logger LOG = LoggerFactory.getLogger(ZNRecordStreamingSerializer.class);

	private static int getListFieldBound(ZNRecord record) {
	int max = Integer.MAX_VALUE;
	if (record.getSimpleFields().containsKey(ZNRecord.LIST_FIELD_BOUND)) {
	String maxStr = record.getSimpleField(ZNRecord.LIST_FIELD_BOUND);
	try {
	max = Integer.parseInt(maxStr);
	} catch (Exception e) {
	LOG.error("IllegalNumberFormat for list field bound: " + maxStr);
	}
	}
	return max;
	}

	@Override
	public byte[] serialize(Object data) throws ZkMarshallingError {
	if (!(data instanceof ZNRecord)) {
	// null is NOT an instance of any class
	LOG.error("Input object must be of type ZNRecord but it is " + data
	+ ". Will not write to zk");
	throw new HelixException("Input object is not of type ZNRecord (was " + data + ")");
	}

	// apply retention policy on list field
	ZNRecord record = (ZNRecord) data;
	int max = getListFieldBound(record);
	if (max < Integer.MAX_VALUE) {
	Map<String, List<String>> listMap = record.getListFields();
	for (String key : listMap.keySet()) {
	List<String> list = listMap.get(key);
	if (list.size() > max) {
	listMap.put(key, list.subList(0, max));
	}
	}
	}
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	byte[] serializedBytes = null;
	try {
	JsonFactory f = new JsonFactory();
	JsonGenerator g = f.createJsonGenerator(baos);

	g.writeStartObject();

	// write id field
	g.writeRaw("\n ");
	g.writeStringField("id", record.getId());

	// write simepleFields
	g.writeRaw("\n ");
	g.writeObjectFieldStart("simpleFields");
	for (String key : record.getSimpleFields().keySet()) {
	g.writeRaw("\n ");
	g.writeStringField(key, record.getSimpleField(key));
	}
	g.writeRaw("\n ");
	g.writeEndObject(); // for simpleFields

	// write listFields
	g.writeRaw("\n ");
	g.writeObjectFieldStart("listFields");
	for (String key : record.getListFields().keySet()) {
	// g.writeStringField(key, record.getListField(key).toString());

	// g.writeObjectFieldStart(key);
	g.writeRaw("\n ");
	g.writeArrayFieldStart(key);
	List<String> list = record.getListField(key);
	for (String listValue : list) {
	g.writeString(listValue);
	}
	// g.writeEndObject();
	g.writeEndArray();

	}
	g.writeRaw("\n ");
	g.writeEndObject(); // for listFields

	// write mapFields
	g.writeRaw("\n ");
	g.writeObjectFieldStart("mapFields");
	for (String key : record.getMapFields().keySet()) {
	// g.writeStringField(key, record.getMapField(key).toString());
	g.writeRaw("\n ");
	g.writeObjectFieldStart(key);
	Map<String, String> map = record.getMapField(key);
	for (String mapKey : map.keySet()) {
	g.writeRaw("\n ");
	g.writeStringField(mapKey, map.get(mapKey));
	}
	g.writeRaw("\n ");
	g.writeEndObject();

	}
	g.writeRaw("\n ");
	g.writeEndObject(); // for mapFields

	byte[] rawPayload = record.getRawPayload();
	if (rawPayload != null && rawPayload.length > 0) {
	// write rawPayload
	g.writeRaw("\n ");
	g.writeStringField("rawPayload", new String(Base64.encodeBase64(rawPayload), "UTF-8"));
	}

	g.writeRaw("\n");
	g.writeEndObject(); // for whole znrecord

	// important: will force flushing of output, close underlying output
	// stream
	g.close();
	serializedBytes = baos.toByteArray();
	// apply compression if needed
	if (record.getBooleanField("enableCompression", false) \|\| serializedBytes.length > ZNRecord.SIZE_LIMIT) {
	serializedBytes = GZipCompressionUtil.compress(serializedBytes);
	}
	} catch (Exception e) {
	LOG.error("Exception during data serialization. Will not write to zk. Data (first 1k): "
	+ new String(baos.toByteArray()).substring(0, 1024), e);
	throw new HelixException(e);
	}
	// check size
	if (serializedBytes.length > ZNRecord.SIZE_LIMIT) {
	LOG.error("Data size larger than 1M, ZNRecord.id: " + record.getId()
	+ ". Will not write to zk. Data (first 1k): "
	+ new String(serializedBytes).substring(0, 1024));
	throw new HelixException("Data size larger than 1M, ZNRecord.id: " + record.getId());
	}

	return serializedBytes;
	}

	@Override
	public Object deserialize(byte[] bytes) throws ZkMarshallingError {
	if (bytes == null \|\| bytes.length == 0) {
	LOG.error("ZNode is empty.");
	return null;
	}

	ByteArrayInputStream bais = new ByteArrayInputStream(bytes);

	ZNRecord record = null;
	String id = null;
	Map<String, String> simpleFields = Maps.newHashMap();
	Map<String, List<String>> listFields = Maps.newHashMap();
	Map<String, Map<String, String>> mapFields = Maps.newHashMap();
	byte[] rawPayload = null;

	try {
	// decompress the data if its already compressed
	if (GZipCompressionUtil.isCompressed(bytes)) {
	byte[] uncompressedBytes = GZipCompressionUtil.uncompress(bais);
	bais = new ByteArrayInputStream(uncompressedBytes);
	}
	JsonFactory f = new JsonFactory();
	JsonParser jp = f.createJsonParser(bais);

	jp.nextToken(); // will return JsonToken.START_OBJECT (verify?)
	while (jp.nextToken() != JsonToken.END_OBJECT) {
	String fieldname = jp.getCurrentName();
	jp.nextToken(); // move to value, or START_OBJECT/START_ARRAY
	if ("id".equals(fieldname)) {
	// contains an object
	id = jp.getText();
	} else if ("simpleFields".equals(fieldname)) {
	while (jp.nextToken() != JsonToken.END_OBJECT) {
	String key = jp.getCurrentName();
	jp.nextToken(); // move to value
	simpleFields.put(key, jp.getText());
	}
	} else if ("mapFields".equals(fieldname)) {
	// user.setVerified(jp.getCurrentToken() == JsonToken.VALUE_TRUE);
	while (jp.nextToken() != JsonToken.END_OBJECT) {
	String key = jp.getCurrentName();
	mapFields.put(key, new TreeMap<String, String>());
	jp.nextToken(); // move to value

	while (jp.nextToken() != JsonToken.END_OBJECT) {
	String mapKey = jp.getCurrentName();
	jp.nextToken(); // move to value
	mapFields.get(key).put(mapKey, jp.getText());
	}
	}

	} else if ("listFields".equals(fieldname)) {
	// user.setUserImage(jp.getBinaryValue());
	while (jp.nextToken() != JsonToken.END_OBJECT) {
	String key = jp.getCurrentName();
	listFields.put(key, new ArrayList<String>());
	jp.nextToken(); // move to value
	while (jp.nextToken() != JsonToken.END_ARRAY) {
	listFields.get(key).add(jp.getText());
	}

	}

	} else if ("rawPayload".equals(fieldname)) {
	rawPayload = Base64.decodeBase64(jp.getText());
	} else {
	throw new IllegalStateException("Unrecognized field '" + fieldname + "'!");
	}
	}
	jp.close(); // ensure resources get cleaned up timely and properly

	if (id == null) {
	throw new IllegalStateException("ZNRecord id field is required!");
	}
	record = new ZNRecord(id);
	record.setSimpleFields(simpleFields);
	record.setListFields(listFields);
	record.setMapFields(mapFields);
	record.setRawPayload(rawPayload);
	} catch (Exception e) {
	LOG.error("Exception during deserialization of bytes: " + new String(bytes), e);
	}
	return record;
	}

	public static void main(String[] args) {
	ZNRecord record = new ZNRecord("record");
	final int recordSize = 10;
	for (int i = 0; i < recordSize; i++) {
	record.setSimpleField("" + i, "" + i);
	record.setListField("" + i, new ArrayList<String>());
	for (int j = 0; j < recordSize; j++) {
	record.getListField("" + i).add("" + j);
	}

	record.setMapField("" + i, new TreeMap<String, String>());
	for (int j = 0; j < recordSize; j++) {
	record.getMapField("" + i).put("" + j, "" + j);
	}
	}

	ZNRecordStreamingSerializer serializer = new ZNRecordStreamingSerializer();
	byte[] bytes = serializer.serialize(record);
	System.out.println(new String(bytes));
	ZNRecord record2 = (ZNRecord) serializer.deserialize(bytes);
	System.out.println(record2);

	long start = System.currentTimeMillis();
	for (int i = 0; i < 100; i++) {
	bytes = serializer.serialize(record);
	// System.out.println(new String(bytes));
	record2 = (ZNRecord) serializer.deserialize(bytes);
	// System.out.println(record2);
	}
	long end = System.currentTimeMillis();
	System.out.println("ZNRecordStreamingSerializer time used: " + (end - start));

	ZNRecordSerializer serializer2 = new ZNRecordSerializer();
	bytes = serializer2.serialize(record);
	// System.out.println(new String(bytes));
	record2 = (ZNRecord) serializer2.deserialize(bytes);
	// System.out.println(record2);

	start = System.currentTimeMillis();
	for (int i = 0; i < 100; i++) {
	bytes = serializer2.serialize(record);
	// System.out.println(new String(bytes));
	record2 = (ZNRecord) serializer2.deserialize(bytes);
	// System.out.println(record2);
	}
	end = System.currentTimeMillis();
	System.out.println("ZNRecordSerializer time used: " + (end - start));

	}
	}