org.apache.servicecomb.foundation.vertx.stream.PumpFromPart.java Source code

Introduction

Here is the source code for org.apache.servicecomb.foundation.vertx.stream.PumpFromPart.java

Source

/*
 * 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.
 */
package org.apache.servicecomb.foundation.vertx.stream;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.concurrent.CompletableFuture;

import javax.servlet.http.Part;

import org.apache.commons.io.IOUtils;
import org.apache.servicecomb.foundation.vertx.http.DownloadUtils;
import org.apache.servicecomb.foundation.vertx.http.ReadStreamPart;

import io.vertx.core.Context;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.streams.ReadStream;
import io.vertx.core.streams.WriteStream;

public class PumpFromPart {
    private Context context;

    private Part part;

    public PumpFromPart(Context context, Part part) {
        this.context = context;
        this.part = part;
    }

    private CompletableFuture<ReadStream<Buffer>> prepareReadStream() {
        CompletableFuture<ReadStream<Buffer>> future = new CompletableFuture<>();

        if (ReadStreamPart.class.isInstance(part)) {
            future.complete(((ReadStreamPart) part).getReadStream());
            return future;
        }

        try {
            InputStream inputStream = part.getInputStream();
            InputStreamToReadStream inputStreamToReadStream = new InputStreamToReadStream(context, inputStream,
                    true);
            inputStreamToReadStream.pause();
            future.complete(inputStreamToReadStream);
        } catch (IOException e) {
            future.completeExceptionally(e);
        }

        return future;
    }

    public CompletableFuture<Void> toWriteStream(WriteStream<Buffer> writeStream) {
        return prepareReadStream()
                .thenCompose(readStream -> new PumpCommon().pump(context, readStream, writeStream))
                .whenComplete((v, e) -> DownloadUtils.clearPartResource(part));
    }

    public CompletableFuture<Void> toOutputStream(OutputStream outputStream, boolean autoCloseOutputStream) {
        if (context == null) {
            return toOutputStreamSync(outputStream, autoCloseOutputStream);
        }

        return toOutputStreamAsync(outputStream, autoCloseOutputStream);
    }

    private CompletableFuture<Void> toOutputStreamAsync(OutputStream outputStream, boolean autoCloseOutputStream) {
        OutputStreamToWriteStream outputStreamToWriteStream = new OutputStreamToWriteStream(context, outputStream,
                autoCloseOutputStream);
        return toWriteStream(outputStreamToWriteStream);
    }

    // DO NOT use a mocked sync context to unify the pump logic
    // otherwise when pump big stream, will cause stack overflow
    private CompletableFuture<Void> toOutputStreamSync(OutputStream outputStream, boolean autoCloseOutputStream) {
        CompletableFuture<Void> future = new CompletableFuture<>();

        try (InputStream inputStream = part.getInputStream()) {
            IOUtils.copyLarge(inputStream, outputStream);
        } catch (Throwable e) {
            future.completeExceptionally(e);
        }

        if (autoCloseOutputStream) {
            try {
                outputStream.close();
            } catch (Throwable e) {
                future.completeExceptionally(e);
            }
        }

        future.complete(null);
        return future;
    }
}