com.yahoo.ycsb.utils.connection.S3Connection.java Source code

Java tutorial

Introduction

Here is the source code for com.yahoo.ycsb.utils.connection.S3Connection.java

Source

/**
 * Copyright (c) 2015 YCSB contributors. All rights reserved.
 * <p>
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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. See accompanying
 * LICENSE file.
 * <p>
 * S3 storage client binding for YCSB.
 */
package com.yahoo.ycsb.utils.connection;

import com.amazonaws.ClientConfiguration;
import com.amazonaws.Protocol;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.regions.Region;
import com.amazonaws.regions.Regions;
import com.amazonaws.services.s3.AmazonS3Client;
import com.amazonaws.services.s3.model.*;
import com.yahoo.ycsb.ClientException;
import com.yahoo.ycsb.Status;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;

// Assumption: there is one bucket per AWS region!

/**
 * S3 Storage client for YCSB framework.
 * <p>
 * Properties to set:
 * <p>
 * s3.accessKeyId=access key S3 aws
 * s3.secretKey=secret key S3 aws
 * s3.endPoint=s3.amazonaws.com
 * s3.region=us-east-1
 * The parameter table is the name of the Bucket where to upload the files.
 * This must be created before starting the benchmark
 * The size of the file to upload is determined by two parameters:
 * - fieldcount this is the number of fields of a record in YCSB
 * - fieldlength this is the size in bytes of a single field in the record
 * together these two parameters define the size of the file to upload,
 * the size in bytes is given by the fieldlength multiplied by the fieldcount.
 * The name of the file is determined by the parameter key.
 * This key is automatically generated by YCSB.
 */
public class S3Connection extends BackendConnection {
    private static Logger logger = Logger.getLogger(S3Connection.class);

    private static String sse = "false";
    private static SSECustomerKey ssecKey;
    private static String accessKeyId = "AKIAJVYH2P5WGTGXVR2Q";
    private static String secretKey = "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";
    private static String maxErrorRetry = "10";
    private static String maxConnections;
    private static String protocol = "HTTP";
    //private static boolean init = true;

    private AmazonS3Client awsClient;
    private String bucket;
    private String region;

    public S3Connection(String bucket, String region, String endPoint) throws ClientException {
        super(bucket, region, endPoint);
        logger.debug("S3Client.establishConnection(" + region + "," + endPoint + ") bucket: " + bucket);
        org.apache.log4j.Logger.getLogger("com.amazonaws").setLevel(Level.OFF);

        /*if (S3Connection.init == true) {
        init();
        S3Connection.init = false;
        }*/

        this.bucket = bucket;
        this.region = region;

        try {
            BasicAWSCredentials s3Credentials = new BasicAWSCredentials(accessKeyId, secretKey);
            ClientConfiguration clientConfig = new ClientConfiguration();
            clientConfig.setMaxErrorRetry(Integer.parseInt(maxErrorRetry));
            if (protocol.equals("HTTP")) {
                clientConfig.setProtocol(Protocol.HTTP);
            } else {
                clientConfig.setProtocol(Protocol.HTTPS);
            }
            if (maxConnections != null) {
                clientConfig.setMaxConnections(Integer.parseInt(maxConnections));
            }

            logger.debug("Inizializing the S3 connection...");
            awsClient = new AmazonS3Client(s3Credentials, clientConfig);
            awsClient.setRegion(Region.getRegion(Regions.fromName(region)));
            awsClient.setEndpoint(endPoint);
            logger.debug("Connection successfully initialized");
        } catch (Exception e) {
            logger.error("Could not connect to S3 storage: " + e.toString());
            e.printStackTrace();
            throw new ClientException(e);
        }
    }

    public String getRegion() {
        return region;
    }

    /**
     * Clean up any state for this storage.
     * Called once per S3 instance;
     */
    public void cleanup() throws ClientException {
        try {
            awsClient.shutdown();
            logger.debug("The client is shutdown successfully");
        } catch (Exception e) {
            logger.error("Could not shutdown the S3Client: " + e.toString());
            e.printStackTrace();
        } finally {
            if (awsClient != null) {
                awsClient = null;
            }
        }
    }

    public Status delete(String key) {
        try {
            awsClient.deleteObject(new DeleteObjectRequest(bucket, key));
        } catch (Exception e) {
            logger.error("Not possible to delete the key " + key);
            e.printStackTrace();
            return Status.ERROR;
        }
        return Status.OK;
    }

    public Status insert(String key, byte[] bytes) {
        try (InputStream input = new ByteArrayInputStream(bytes)) {
            ObjectMetadata metadata = new ObjectMetadata();
            metadata.setContentLength(bytes.length);
            PutObjectRequest putObjectRequest = null;
            if (ssecKey != null) {
                if (ssecKey.equals("true")) {
                    metadata.setSSEAlgorithm(ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION);
                    putObjectRequest = new PutObjectRequest(bucket, key, input, metadata);
                } else {
                    putObjectRequest = new PutObjectRequest(bucket, key, input, metadata)
                            .withSSECustomerKey(ssecKey);
                }
            } else {
                putObjectRequest = new PutObjectRequest(bucket, key, input, metadata);
            }

            try {
                PutObjectResult res = awsClient.putObject(putObjectRequest);
                if (res.getETag() == null) {
                    return Status.ERROR;
                } else {
                    if (ssecKey != null) {
                        if (ssecKey.equals("true")) {
                            logger.debug("Uploaded object encryption status is " + res.getSSEAlgorithm());
                        } else {
                            logger.debug("Uploaded object encryption status is " + res.getSSEAlgorithm());
                        }
                    }
                }
            } catch (Exception e) {
                logger.error("Not possible to write object :" + key);
                System.err.println("Retrying " + key);
                insert(key, bytes);
            }
        } catch (Exception e) {
            logger.error("Error in the creation of the stream :" + e.toString());
            System.err.println("Retrying " + key);
            insert(key, bytes);
            //e.printStackTrace();
            //return Status.ERROR;
        }
        return Status.OK;
    }

    public byte[] read(String key) throws InterruptedException {
        //long starttime = System.currentTimeMillis();
        byte[] bytes = null;
        try {
            GetObjectRequest getObjectRequest = null;
            GetObjectMetadataRequest getObjectMetadataRequest = null;
            if (ssecKey != null) {
                getObjectRequest = new GetObjectRequest(bucket, key).withSSECustomerKey(ssecKey);
                getObjectMetadataRequest = new GetObjectMetadataRequest(bucket, key).withSSECustomerKey(ssecKey);
            } else {
                getObjectRequest = new GetObjectRequest(bucket, key);
                getObjectMetadataRequest = new GetObjectMetadataRequest(bucket, key);
            }
            //System.out.println("Get object " + key + " from " + bucket);
            S3Object object = awsClient.getObject(getObjectRequest);
            ObjectMetadata objectMetadata = awsClient.getObjectMetadata(getObjectMetadataRequest);
            //System.out.println("Get object " + key + " from " + bucket + " OK");
            InputStream objectData = object.getObjectContent(); //consuming the stream
            // writing the stream to bytes and to results
            int sizeOfFile = (int) objectMetadata.getContentLength();
            bytes = new byte[sizeOfFile];
            int offset = 0;
            while (offset < sizeOfFile) {
                int chunk_size;

                //read in 4k chunks
                chunk_size = sizeOfFile - offset > 4096 ? 4096 : sizeOfFile - offset;

                int nr_bytes_read = objectData.read(bytes, offset, sizeOfFile - offset);
                offset = offset + nr_bytes_read;

                if (Thread.interrupted()) {
                    //System.out.println("interrupt " + key);
                    objectData.close();
                    throw new InterruptedException();
                }
            }
            //int nr_bytes_read = objectData.read(bytes, 0, sizeOfFile);
            objectData.close();
        } catch (IOException e) {
            logger.warn("Not possible to get the object " + key);
        }
        //long endtime = System.currentTimeMillis();
        //System.out.println("ReadS3: " + key + " " + (endtime - starttime) + " " + region);
        return bytes;
    }
}