co.cask.cdap.data.stream.service.upload.BufferedContentWriter.java Source code

Introduction

Here is the source code for co.cask.cdap.data.stream.service.upload.BufferedContentWriter.java

Source

/*
 * Copyright  2015 Cask Data, Inc.
 *
 * Licensed 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.
 */

package co.cask.cdap.data.stream.service.upload;

import co.cask.cdap.api.stream.StreamEventData;
import co.cask.cdap.common.NotFoundException;
import co.cask.cdap.common.io.ByteBuffers;
import co.cask.cdap.data.stream.service.ConcurrentStreamWriter;
import co.cask.cdap.data.stream.service.MutableStreamEventData;
import co.cask.cdap.proto.Id;
import com.google.common.base.Throwables;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * A {@link ContentWriter} that buffer all events in memory and write to the actual stream writer when closed.
 */
final class BufferedContentWriter implements ContentWriter, Iterable<ByteBuffer> {

    private final Id.Stream streamId;
    private final ConcurrentStreamWriter streamWriter;
    private final Map<String, String> headers;
    private final List<ByteBuffer> bodies;

    BufferedContentWriter(Id.Stream streamId, ConcurrentStreamWriter streamWriter, Map<String, String> headers) {
        this.streamId = streamId;
        this.streamWriter = streamWriter;
        this.headers = ImmutableMap.copyOf(headers);
        this.bodies = Lists.newLinkedList();
    }

    @Override
    public void append(ByteBuffer body, boolean immutable) throws IOException {
        if (immutable) {
            bodies.add(body);
        } else {
            bodies.add(ByteBuffers.copy(body));
        }
    }

    @Override
    public void appendAll(Iterator<ByteBuffer> bodies, boolean immutable) throws IOException {
        while (bodies.hasNext()) {
            append(bodies.next(), immutable);
        }
    }

    @Override
    public void close() throws IOException {
        try {
            streamWriter.enqueue(streamId, new StreamEventDataIterator(headers, bodies.iterator()));
        } catch (NotFoundException e) {
            throw Throwables.propagate(e);
        }
    }

    @Override
    public void cancel() {
        // No-op
    }

    @Override
    public Iterator<ByteBuffer> iterator() {
        return bodies.iterator();
    }

    private static final class StreamEventDataIterator extends AbstractIterator<StreamEventData> {

        private final Iterator<? extends ByteBuffer> bodies;
        private final MutableStreamEventData streamEventData;

        private StreamEventDataIterator(Map<String, String> headers, Iterator<? extends ByteBuffer> bodies) {
            this.bodies = bodies;
            this.streamEventData = new MutableStreamEventData().setHeaders(headers);
        }

        @Override
        protected StreamEventData computeNext() {
            if (!bodies.hasNext()) {
                return endOfData();
            }
            return streamEventData.setBody(bodies.next());
        }
    }
}