Andrii Shvaika
5 years ago
4 changed files with 237 additions and 4 deletions
@ -0,0 +1,185 @@ |
|||
/* |
|||
* 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. |
|||
*/ |
|||
|
|||
/* |
|||
* Content of this file was modified to addresses the issue https://issues.apache.org/jira/browse/KAFKA-4090
|
|||
* |
|||
*/ |
|||
package org.apache.kafka.common.network; |
|||
|
|||
import org.apache.kafka.common.memory.MemoryPool; |
|||
import org.slf4j.Logger; |
|||
import org.slf4j.LoggerFactory; |
|||
import org.thingsboard.server.common.data.exception.ThingsboardKafkaClientError; |
|||
|
|||
import java.io.EOFException; |
|||
import java.io.IOException; |
|||
import java.nio.ByteBuffer; |
|||
import java.nio.channels.ScatteringByteChannel; |
|||
import java.util.Arrays; |
|||
import java.util.stream.Collectors; |
|||
|
|||
/** |
|||
* A size delimited Receive that consists of a 4 byte network-ordered size N followed by N bytes of content |
|||
*/ |
|||
public class NetworkReceive implements Receive { |
|||
|
|||
public final static String UNKNOWN_SOURCE = ""; |
|||
public final static int UNLIMITED = -1; |
|||
public final static int TB_MAX_REQUESTED_BUFFER_SIZE = 100 * 1024 * 1024; |
|||
public final static int TB_LOG_REQUESTED_BUFFER_SIZE = 10 * 1024 * 1024; |
|||
private static final Logger log = LoggerFactory.getLogger(NetworkReceive.class); |
|||
private static final ByteBuffer EMPTY_BUFFER = ByteBuffer.allocate(0); |
|||
|
|||
private final String source; |
|||
private final ByteBuffer size; |
|||
private final int maxSize; |
|||
private final MemoryPool memoryPool; |
|||
private int requestedBufferSize = -1; |
|||
private ByteBuffer buffer; |
|||
|
|||
|
|||
public NetworkReceive(String source, ByteBuffer buffer) { |
|||
this.source = source; |
|||
this.buffer = buffer; |
|||
this.size = null; |
|||
this.maxSize = TB_MAX_REQUESTED_BUFFER_SIZE; |
|||
this.memoryPool = MemoryPool.NONE; |
|||
} |
|||
|
|||
public NetworkReceive(String source) { |
|||
this.source = source; |
|||
this.size = ByteBuffer.allocate(4); |
|||
this.buffer = null; |
|||
this.maxSize = TB_MAX_REQUESTED_BUFFER_SIZE; |
|||
this.memoryPool = MemoryPool.NONE; |
|||
} |
|||
|
|||
public NetworkReceive(int maxSize, String source) { |
|||
this.source = source; |
|||
this.size = ByteBuffer.allocate(4); |
|||
this.buffer = null; |
|||
this.maxSize = getMaxSize(maxSize); |
|||
this.memoryPool = MemoryPool.NONE; |
|||
} |
|||
|
|||
public NetworkReceive(int maxSize, String source, MemoryPool memoryPool) { |
|||
this.source = source; |
|||
this.size = ByteBuffer.allocate(4); |
|||
this.buffer = null; |
|||
this.maxSize = getMaxSize(maxSize); |
|||
this.memoryPool = memoryPool; |
|||
} |
|||
|
|||
public NetworkReceive() { |
|||
this(UNKNOWN_SOURCE); |
|||
} |
|||
|
|||
@Override |
|||
public String source() { |
|||
return source; |
|||
} |
|||
|
|||
@Override |
|||
public boolean complete() { |
|||
return !size.hasRemaining() && buffer != null && !buffer.hasRemaining(); |
|||
} |
|||
|
|||
public long readFrom(ScatteringByteChannel channel) throws IOException { |
|||
int read = 0; |
|||
if (size.hasRemaining()) { |
|||
int bytesRead = channel.read(size); |
|||
if (bytesRead < 0) |
|||
throw new EOFException(); |
|||
read += bytesRead; |
|||
if (!size.hasRemaining()) { |
|||
size.rewind(); |
|||
int receiveSize = size.getInt(); |
|||
if (receiveSize < 0) |
|||
throw new InvalidReceiveException("Invalid receive (size = " + receiveSize + ")"); |
|||
if (maxSize != UNLIMITED && receiveSize > maxSize) { |
|||
var t = Thread.currentThread(); |
|||
throw new ThingsboardKafkaClientError("Invalid receive (size = " + receiveSize + " larger than " + maxSize + ")"); |
|||
} |
|||
requestedBufferSize = receiveSize; //may be 0 for some payloads (SASL)
|
|||
if (receiveSize == 0) { |
|||
buffer = EMPTY_BUFFER; |
|||
} |
|||
} |
|||
} |
|||
if (buffer == null && requestedBufferSize != -1) { //we know the size we want but havent been able to allocate it yet
|
|||
if (requestedBufferSize > TB_LOG_REQUESTED_BUFFER_SIZE) { |
|||
String stackTrace = Arrays.stream(Thread.currentThread().getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("|")); |
|||
log.error("Allocating buffer of size {} for source {}", requestedBufferSize, source); |
|||
log.error("Stack Trace: {}", stackTrace); |
|||
} |
|||
buffer = memoryPool.tryAllocate(requestedBufferSize); |
|||
if (buffer == null) |
|||
log.trace("Broker low on memory - could not allocate buffer of size {} for source {}", requestedBufferSize, source); |
|||
} |
|||
if (buffer != null) { |
|||
int bytesRead = channel.read(buffer); |
|||
if (bytesRead < 0) |
|||
throw new EOFException(); |
|||
read += bytesRead; |
|||
} |
|||
|
|||
return read; |
|||
} |
|||
|
|||
@Override |
|||
public boolean requiredMemoryAmountKnown() { |
|||
return requestedBufferSize != -1; |
|||
} |
|||
|
|||
@Override |
|||
public boolean memoryAllocated() { |
|||
return buffer != null; |
|||
} |
|||
|
|||
|
|||
@Override |
|||
public void close() throws IOException { |
|||
if (buffer != null && buffer != EMPTY_BUFFER) { |
|||
memoryPool.release(buffer); |
|||
buffer = null; |
|||
} |
|||
} |
|||
|
|||
public ByteBuffer payload() { |
|||
return this.buffer; |
|||
} |
|||
|
|||
public int bytesRead() { |
|||
if (buffer == null) |
|||
return size.position(); |
|||
return buffer.position() + size.position(); |
|||
} |
|||
|
|||
/** |
|||
* Returns the total size of the receive including payload and size buffer |
|||
* for use in metrics. This is consistent with {@link NetworkSend#size()} |
|||
*/ |
|||
public int size() { |
|||
return payload().limit() + size.limit(); |
|||
} |
|||
|
|||
private int getMaxSize(int maxSize) { |
|||
return maxSize == UNLIMITED ? TB_MAX_REQUESTED_BUFFER_SIZE : Math.min(maxSize, TB_MAX_REQUESTED_BUFFER_SIZE); |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,23 @@ |
|||
/** |
|||
* Copyright © 2016-2022 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.common.data.exception; |
|||
|
|||
public class ThingsboardKafkaClientError extends Error { |
|||
|
|||
public ThingsboardKafkaClientError(String message) { |
|||
super(message); |
|||
} |
|||
} |
|||
Loading…
Reference in new issue