spark/network/TestManagedBuffer.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.network;
0019
0020 import java.io.IOException;
0021 import java.io.InputStream;
0022 import java.nio.ByteBuffer;
0023
0024 import com.google.common.base.Preconditions;
0025 import io.netty.buffer.Unpooled;
0026
0027 import org.apache.spark.network.buffer.ManagedBuffer;
0028 import org.apache.spark.network.buffer.NettyManagedBuffer;
0029
0030 /**
0031  * A ManagedBuffer implementation that contains 0, 1, 2, 3, ..., (len-1).
0032  *
0033  * Used for testing.
0034  */
0035 public class TestManagedBuffer extends ManagedBuffer {
0036
0037   private final int len;
0038   private NettyManagedBuffer underlying;
0039
0040   public TestManagedBuffer(int len) {
0041     Preconditions.checkArgument(len <= Byte.MAX_VALUE);
0042     this.len = len;
0043     byte[] byteArray = new byte[len];
0044     for (int i = 0; i < len; i ++) {
0045       byteArray[i] = (byte) i;
0046     }
0047     this.underlying = new NettyManagedBuffer(Unpooled.wrappedBuffer(byteArray));
0048   }
0049
0050
0051   @Override
0052   public long size() {
0053     return underlying.size();
0054   }
0055
0056   @Override
0057   public ByteBuffer nioByteBuffer() throws IOException {
0058     return underlying.nioByteBuffer();
0059   }
0060
0061   @Override
0062   public InputStream createInputStream() throws IOException {
0063     return underlying.createInputStream();
0064   }
0065
0066   @Override
0067   public ManagedBuffer retain() {
0068     underlying.retain();
0069     return this;
0070   }
0071
0072   @Override
0073   public ManagedBuffer release() {
0074     underlying.release();
0075     return this;
0076   }
0077
0078   @Override
0079   public Object convertToNetty() throws IOException {
0080     return underlying.convertToNetty();
0081   }
0082
0083   @Override
0084   public int hashCode() {
0085     return underlying.hashCode();
0086   }
0087
0088   @Override
0089   public boolean equals(Object other) {
0090     if (other instanceof ManagedBuffer) {
0091       try {
0092         ByteBuffer nioBuf = ((ManagedBuffer) other).nioByteBuffer();
0093         if (nioBuf.remaining() != len) {
0094           return false;
0095         } else {
0096           for (int i = 0; i < len; i ++) {
0097             if (nioBuf.get() != i) {
0098               return false;
0099             }
0100           }
0101           return true;
0102         }
0103       } catch (IOException e) {
0104         throw new RuntimeException(e);
0105       }
0106     }
0107     return false;
0108   }
0109 }