/*
* JBoss, Home of Professional Open Source.
* Copyright 2014 Red Hat, Inc., and individual contributors
* as indicated by the @author tags.
*
* 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 io.undertow.util;
import io.undertow.
UndertowMessages;
import io.undertow.connector.
PooledByteBuffer;
import java.nio.
ByteBuffer;
import java.util.concurrent.atomic.
AtomicIntegerFieldUpdater;
/**
* A reference counted pooled implementation, that basically consists of a main buffer, that can be sliced off into smaller buffers,
* and the underlying buffer will not be freed until all the slices and the main buffer itself have also been freed.
*
* This also supports the notion of un-freeing the main buffer. Basically this allows the buffer be re-used, so if only a small slice of the
* buffer was used for read operations the main buffer can potentially be re-used. This prevents buffer exhaustion attacks where content
* is sent in many small packets, and you end up allocating a large number of buffers to hold a small amount of data.
*
* @author Stuart Douglas
*/
public class
ReferenceCountedPooled implements
PooledByteBuffer {
private final
PooledByteBuffer underlying;
@
SuppressWarnings("unused")
private volatile int
referenceCount;
boolean
mainFreed = false;
private
ByteBuffer slice = null;
private final
FreeNotifier freeNotifier;
private static final
AtomicIntegerFieldUpdater<
ReferenceCountedPooled>
referenceCountUpdater =
AtomicIntegerFieldUpdater.
newUpdater(
ReferenceCountedPooled.class, "referenceCount");
public
ReferenceCountedPooled(
PooledByteBuffer underlying, int
referenceCount) {
this(
underlying,
referenceCount, null);
}
public
ReferenceCountedPooled(
PooledByteBuffer underlying, int
referenceCount,
FreeNotifier freeNotifier) {
this.
underlying =
underlying;
this.
referenceCount =
referenceCount;
this.
freeNotifier =
freeNotifier;
}
@
Override
public void
close() {
if(
mainFreed) {
return;
}
mainFreed = true;
freeInternal();
}
@
Override
public boolean
isOpen() {
return !
mainFreed;
}
public boolean
isFreed() {
return
mainFreed;
}
public boolean
tryUnfree() {
int
refs;
do {
refs =
referenceCountUpdater.
get(this);
if(
refs <= 0) {
return false;
}
} while (!
referenceCountUpdater.
compareAndSet(this,
refs,
refs + 1));
ByteBuffer resource =
slice != null ?
slice :
underlying.
getBuffer();
resource.
position(
resource.
limit());
resource.
limit(
resource.
capacity());
slice =
resource.
slice();
mainFreed = false;
return true;
}
private void
freeInternal() {
if(
referenceCountUpdater.
decrementAndGet(this) == 0) {
underlying.
close();
if(
freeNotifier != null) {
freeNotifier.
freed();
}
}
}
@
Override
public
ByteBuffer getBuffer() throws
IllegalStateException {
if(
mainFreed) {
throw
UndertowMessages.
MESSAGES.
bufferAlreadyFreed();
}
if(
slice != null) {
return
slice;
}
return
underlying.
getBuffer();
}
public
PooledByteBuffer createView(final
ByteBuffer newValue) {
increaseReferenceCount();
return new
PooledByteBuffer() {
boolean
free = false;
@
Override
public void
close() {
//make sure that a given view can only be freed once
if(!
free) {
free = true;
ReferenceCountedPooled.this.
freeInternal();
}
}
@
Override
public boolean
isOpen() {
return !
free;
}
@
Override
public
ByteBuffer getBuffer() throws
IllegalStateException {
if(
free) {
throw
UndertowMessages.
MESSAGES.
bufferAlreadyFreed();
}
return
newValue;
}
@
Override
public
String toString() {
return "ReferenceCountedPooled$view{" +
"buffer=" +
newValue +
"free=" +
free +
"underlying=" +
underlying +
", referenceCount=" +
referenceCount +
", mainFreed=" +
mainFreed +
", slice=" +
slice +
'}';
}
};
}
public void
increaseReferenceCount() {
int
val;
do {
val =
referenceCountUpdater.
get(this);
if(
val == 0) {
//should never happen, as this should only be called from
//code that already has a reference
throw
UndertowMessages.
MESSAGES.
objectWasFreed();
}
} while (!
referenceCountUpdater.
compareAndSet(this,
val,
val + 1));
}
public interface
FreeNotifier {
void
freed();
}
@
Override
public
String toString() {
return "ReferenceCountedPooled{" +
"underlying=" +
underlying +
", referenceCount=" +
referenceCount +
", mainFreed=" +
mainFreed +
", slice=" +
slice +
'}';
}
}