Source code

001/* 
002 * This file is part of the Kompics component model runtime.
003 *
004 * Copyright (C) 2009 Swedish Institute of Computer Science (SICS) 
005 * Copyright (C) 2009 Royal Institute of Technology (KTH)
006 *
007 * This program is free software; you can redistribute it and/or
008 * modify it under the terms of the GNU General Public License
009 * as published by the Free Software Foundation; either version 2
010 * of the License, or (at your option) any later version.
011 *
012 * This program is distributed in the hope that it will be useful,
013 * but WITHOUT ANY WARRANTY; without even the implied warranty of
014 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
015 * GNU General Public License for more details.
016 *
017 * You should have received a copy of the GNU General Public License
018 * along with this program; if not, write to the Free Software
019 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
020 */
021package se.sics.kompics.network.netty.serialization;
022
023import java.util.Optional;
024import com.google.common.io.Closer;
025import io.netty.buffer.ByteBuf;
026import io.netty.buffer.ByteBufInputStream;
027import io.netty.buffer.ByteBufOutputStream;
028import java.io.IOException;
029import java.util.concurrent.ConcurrentHashMap;
030import java.util.concurrent.ConcurrentMap;
031import java.util.concurrent.ConcurrentSkipListMap;
032import org.apache.avro.Schema;
033import org.apache.avro.file.DataFileStream;
034import org.apache.avro.file.DataFileWriter;
035import org.apache.avro.generic.GenericContainer;
036import org.apache.avro.io.BinaryDecoder;
037import org.apache.avro.io.BinaryEncoder;
038import org.apache.avro.io.DatumReader;
039import org.apache.avro.io.DatumWriter;
040import org.apache.avro.io.DecoderFactory;
041import org.apache.avro.io.EncoderFactory;
042import org.apache.avro.reflect.ReflectData;
043import org.apache.avro.reflect.ReflectDatumReader;
044import org.apache.avro.reflect.ReflectDatumWriter;
045import org.apache.avro.specific.SpecificDatumReader;
046import org.apache.avro.specific.SpecificDatumWriter;
047import org.slf4j.Logger;
048import org.slf4j.LoggerFactory;
049import se.sics.kompics.network.netty.serialization.SpecialSerializers.BitBuffer;
050
051/**
052 *
053 * @author lkroll
054 */
055public class AvroSerializer implements Serializer {
056
057    private static final Logger LOG = LoggerFactory.getLogger(AvroSerializer.class);
058
059    private static final int MY_ID = 7;
060
061    private static final ConcurrentMap<Integer, SchemaEntry> idMap = new ConcurrentSkipListMap<Integer, SchemaEntry>();
062    private static final ConcurrentMap<String, SchemaEntry> classMap = new ConcurrentHashMap<String, SchemaEntry>();
063
064    private static final ReflectData rData = ReflectData.get();
065
066    public static synchronized void register(int id, Class<?> type) throws KeyExistsException, InvalidKeyException {
067        register(id, type, false);
068    }
069
070    public static synchronized void register(int id, Class<?> type, boolean force)
071            throws KeyExistsException, InvalidKeyException {
072        String typeName = type.getName();
073        if (!force) {
074            // Check first once, since schema generation is expensive
075            if (idMap.containsKey(id)) {
076                throw new KeyExistsException(id);
077            }
078            if (classMap.containsKey(typeName)) {
079                throw new KeyExistsException(typeName);
080            }
081        }
082        if (id < 0) {
083            throw new InvalidKeyException(id);
084        }
085        Schema s = rData.getSchema(type);
086        SchemaEntry se = new SchemaEntry(s, type, id, false);
087        idMap.put(id, se);
088        classMap.put(typeName, se);
089
090        Serializers.register(type, MY_ID);
091    }
092
093    public static synchronized void register(int id, Class<?> type, Schema schema)
094            throws KeyExistsException, InvalidKeyException {
095        register(id, type, schema, false);
096    }
097
098    public static synchronized void register(int id, Class<?> type, Schema schema, boolean force)
099            throws KeyExistsException, InvalidKeyException {
100        String typeName = type.getName();
101        if (!force) {
102            // Check first once, since schema generation is expensive
103            if (idMap.containsKey(id)) {
104                throw new KeyExistsException(id);
105            }
106            if (classMap.containsKey(typeName)) {
107                throw new KeyExistsException(typeName);
108            }
109        }
110        if (id < 0) {
111            throw new InvalidKeyException(id);
112        }
113        SchemaEntry se = new SchemaEntry(schema, type, id, true);
114        idMap.put(id, se);
115        classMap.put(typeName, se);
116
117        Serializers.register(type, MY_ID);
118    }
119
120    @Override
121    public int identifier() {
122        return MY_ID;
123    }
124
125    @Override
126    public void toBinary(Object o, ByteBuf buf) {
127        Class<?> type = o.getClass();
128        String typeName = type.getName();
129        SchemaEntry se = classMap.get(typeName);
130        if (se == null) {
131            toBinaryNoSchema(o, type, buf);
132        } else {
133            toBinaryWithSchema(o, se, buf);
134        }
135    }
136
137    private void toBinaryWithSchema(Object o, SchemaEntry se, ByteBuf buf) {
138        BitBuffer flags;
139        if (se.generated) {
140            flags = BitBuffer.create(true, true); // with schema and generated
141        } else {
142            flags = BitBuffer.create(true, false); // with schema but not generated
143        }
144        byte[] flagsB = flags.finalise();
145        buf.writeBytes(flagsB);
146        buf.writeInt(se.id);
147        try {
148            Closer closer = Closer.create(); // Did I mention how much java6 sucks?
149            try {
150                ByteBufOutputStream out = closer.register(new ByteBufOutputStream(buf));
151                BinaryEncoder encoder = EncoderFactory.get().directBinaryEncoder(out, null);
152                DatumWriter<Object> writer;
153                if (se.generated) {
154                    writer = new SpecificDatumWriter<Object>(se.schema);
155                } else {
156                    writer = new ReflectDatumWriter<Object>(se.schema);
157                }
158                writer.write(o, encoder);
159                encoder.flush();
160            } catch (Throwable ex) {
161                LOG.error("Couldn't serialise object.", ex);
162                closer.rethrow(ex);
163            } finally {
164                closer.close();
165            }
166        } catch (IOException ex) {
167            LOG.error("Couldn't serialise object.", ex);
168        }
169    }
170
171    private void toBinaryNoSchema(Object o, Class<?> type, ByteBuf buf) {
172        LOG.info(
173                "Prepending schema to object of type {}. This is not efficient. It's recommended to register the class instead.",
174                type);
175        Schema s;
176        BitBuffer flags;
177        DatumWriter<Object> refWriter;
178        if (o instanceof GenericContainer) {
179            GenericContainer ag = (GenericContainer) o;
180            s = ag.getSchema();
181            flags = BitBuffer.create(false, true); // no schema and generated
182            refWriter = new SpecificDatumWriter<>(s);
183        } else {
184            s = rData.getSchema(type);
185            flags = BitBuffer.create(false, false); // no schema and not generated
186            refWriter = new ReflectDatumWriter<>(s);
187        }
188        byte[] flagsB = flags.finalise();
189        buf.writeBytes(flagsB);
190        try {
191            Closer closer = Closer.create(); // Did I mention how much java6 sucks?
192            try {
193                ByteBufOutputStream out = closer.register(new ByteBufOutputStream(buf));
194                @SuppressWarnings("resource")
195                DataFileWriter<Object> writer = closer.register(new DataFileWriter<>(refWriter).create(s, out));
196                writer.append(o);
197                writer.flush();
198            } catch (Throwable ex) {
199                LOG.error("Couldn't serialise object.", ex);
200                closer.rethrow(ex);
201            } finally {
202                closer.close();
203            }
204        } catch (IOException ex) {
205            LOG.error("Couldn't serialise object.", ex);
206        }
207    }
208
209    @Override
210    public Object fromBinary(ByteBuf buf, Optional<Object> hint) {
211        byte[] flagsB = new byte[1];
212        buf.readBytes(flagsB);
213        boolean[] flags = BitBuffer.extract(2, flagsB);
214        boolean registered = flags[0];
215        boolean generated = flags[1];
216        if (registered) {
217            int id = buf.readInt();
218            SchemaEntry se = idMap.get(id);
219            if (se == null) {
220                LOG.warn("Could not deserialize object for id {}! Not registered!", id);
221                return null;
222            }
223            return fromBinaryWithSchema(buf, se, generated);
224        } else {
225            return fromBinaryNoSchema(buf, generated);
226        }
227    }
228
229    private Object fromBinaryNoSchema(ByteBuf buf, boolean generated) {
230        DatumReader<Object> refReader;
231        if (generated) {
232            refReader = new SpecificDatumReader<>();
233        } else {
234            refReader = new ReflectDatumReader<>();
235        }
236        try {
237            Closer closer = Closer.create();
238            try {
239                ByteBufInputStream in = closer.register(new ByteBufInputStream(buf));
240                DataFileStream<Object> reader = closer.register(new DataFileStream<>(in, refReader));
241                return reader.next(); // there should only be one
242            } catch (Throwable ex) {
243                LOG.error("Couldn't deserialise object.", ex);
244                closer.rethrow(ex);
245                return null;
246            } finally {
247                closer.close();
248            }
249        } catch (IOException ex) {
250            LOG.error("Couldn't deserialise object.", ex);
251            return null;
252        }
253    }
254
255    private Object fromBinaryWithSchema(ByteBuf buf, SchemaEntry se, boolean generated) {
256        DatumReader<Object> refReader;
257        if (generated) {
258            refReader = new SpecificDatumReader<>(se.schema);
259        } else {
260            refReader = new ReflectDatumReader<>(se.schema);
261        }
262        try {
263            Closer closer = Closer.create();
264            try {
265                ByteBufInputStream in = closer.register(new ByteBufInputStream(buf));
266                BinaryDecoder decoder = DecoderFactory.get().directBinaryDecoder(in, null);
267                return refReader.read(null, decoder); // there should only be one
268            } catch (Throwable ex) {
269                LOG.error("Couldn't deserialise object.", ex);
270                closer.rethrow(ex);
271                return null;
272            } finally {
273                closer.close();
274            }
275        } catch (IOException ex) {
276            LOG.error("Couldn't deserialise object.", ex);
277            return null;
278        }
279    }
280
281    private static class SchemaEntry {
282
283        public final Schema schema;
284        @SuppressWarnings("unused")
285        public final Class<?> type;
286        public final int id;
287        public final boolean generated;
288
289        public SchemaEntry(Schema schema, Class<?> type, int id, boolean generated) {
290            this.schema = schema;
291            this.type = type;
292            this.id = id;
293            this.generated = generated;
294        }
295    }
296
297    @SuppressWarnings("serial")
298    public static class KeyExistsException extends Exception {
299
300        private final Object key;
301
302        public KeyExistsException(Object key) {
303            this.key = key;
304        }
305
306        @Override
307        public String getMessage() {
308            return "Key " + key + " already exists!";
309        }
310    }
311
312    @SuppressWarnings("serial")
313    public static class InvalidKeyException extends Exception {
314
315        private final int key;
316
317        public InvalidKeyException(int key) {
318            this.key = key;
319        }
320
321        @Override
322        public String getMessage() {
323            return "Key " + key + " is invalid! Must be positive integer.";
324        }
325    }
326}