network/client/StreamInterceptor.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.client;
0019
0020 import java.nio.ByteBuffer;
0021 import java.nio.channels.ClosedChannelException;
0022
0023 import io.netty.buffer.ByteBuf;
0024
0025 import org.apache.spark.network.protocol.Message;
0026 import org.apache.spark.network.server.MessageHandler;
0027 import org.apache.spark.network.util.TransportFrameDecoder;
0028
0029 /**
0030  * An interceptor that is registered with the frame decoder to feed stream data to a
0031  * callback.
0032  */
0033 public class StreamInterceptor<T extends Message> implements TransportFrameDecoder.Interceptor {
0034
0035   private final MessageHandler<T> handler;
0036   private final String streamId;
0037   private final long byteCount;
0038   private final StreamCallback callback;
0039   private long bytesRead;
0040
0041   public StreamInterceptor(
0042       MessageHandler<T> handler,
0043       String streamId,
0044       long byteCount,
0045       StreamCallback callback) {
0046     this.handler = handler;
0047     this.streamId = streamId;
0048     this.byteCount = byteCount;
0049     this.callback = callback;
0050     this.bytesRead = 0;
0051   }
0052
0053   @Override
0054   public void exceptionCaught(Throwable cause) throws Exception {
0055     deactivateStream();
0056     callback.onFailure(streamId, cause);
0057   }
0058
0059   @Override
0060   public void channelInactive() throws Exception {
0061     deactivateStream();
0062     callback.onFailure(streamId, new ClosedChannelException());
0063   }
0064
0065   private void deactivateStream() {
0066     if (handler instanceof TransportResponseHandler) {
0067       // we only have to do this for TransportResponseHandler as it exposes numOutstandingFetches
0068       // (there is no extra cleanup that needs to happen)
0069       ((TransportResponseHandler) handler).deactivateStream();
0070     }
0071   }
0072
0073   @Override
0074   public boolean handle(ByteBuf buf) throws Exception {
0075     int toRead = (int) Math.min(buf.readableBytes(), byteCount - bytesRead);
0076     ByteBuffer nioBuffer = buf.readSlice(toRead).nioBuffer();
0077
0078     int available = nioBuffer.remaining();
0079     callback.onData(streamId, nioBuffer);
0080     bytesRead += available;
0081     if (bytesRead > byteCount) {
0082       RuntimeException re = new IllegalStateException(String.format(
0083         "Read too many bytes? Expected %d, but read %d.", byteCount, bytesRead));
0084       callback.onFailure(streamId, re);
0085       deactivateStream();
0086       throw re;
0087     } else if (bytesRead == byteCount) {
0088       deactivateStream();
0089       callback.onComplete(streamId);
0090     }
0091
0092     return bytesRead != byteCount;
0093   }
0094
0095 }