/**
* 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.
*/
package org.apache.hadoop.yarn.server.resourcemanager.reservation;
import java.io.IOException;
import java.io.StringWriter;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;
import java.util.NavigableMap;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.util.resource.ResourceCalculator;
import org.apache.hadoop.yarn.util.resource.Resources;
import com.google.gson.stream.JsonWriter;
/**
* This is a run length encoded sparse data structure that maintains resource
* allocations over time.
*/
public class RLESparseResourceAllocation {
private static final int THRESHOLD = 100;
private static final Resource ZERO_RESOURCE = Resource.newInstance(0, 0);
private TreeMap<Long, Resource> cumulativeCapacity =
new TreeMap<Long, Resource>();
private final ReentrantReadWriteLock readWriteLock =
new ReentrantReadWriteLock();
private final Lock readLock = readWriteLock.readLock();
private final Lock writeLock = readWriteLock.writeLock();
private final ResourceCalculator resourceCalculator;
private final Resource minAlloc;
public RLESparseResourceAllocation(ResourceCalculator resourceCalculator,
Resource minAlloc) {
this.resourceCalculator = resourceCalculator;
this.minAlloc = minAlloc;
}
private boolean isSameAsPrevious(Long key, Resource capacity) {
Entry<Long, Resource> previous = cumulativeCapacity.lowerEntry(key);
return (previous != null && previous.getValue().equals(capacity));
}
private boolean isSameAsNext(Long key, Resource capacity) {
Entry<Long, Resource> next = cumulativeCapacity.higherEntry(key);
return (next != null && next.getValue().equals(capacity));
}
/**
* Add a resource for the specified interval
*
* @param reservationInterval the interval for which the resource is to be
* added
* @param totCap the resource to be added
* @return true if addition is successful, false otherwise
*/
public boolean addInterval(ReservationInterval reservationInterval,
Resource totCap) {
if (totCap.equals(ZERO_RESOURCE)) {
return true;
}
writeLock.lock();
try {
long startKey = reservationInterval.getStartTime();
long endKey = reservationInterval.getEndTime();
NavigableMap<Long, Resource> ticks =
cumulativeCapacity.headMap(endKey, false);
if (ticks != null && !ticks.isEmpty()) {
Resource updatedCapacity = Resource.newInstance(0, 0);
Entry<Long, Resource> lowEntry = ticks.floorEntry(startKey);
if (lowEntry == null) {
// This is the earliest starting interval
cumulativeCapacity.put(startKey, totCap);
} else {
updatedCapacity = Resources.add(lowEntry.getValue(), totCap);
// Add a new tick only if the updated value is different
// from the previous tick
if ((startKey == lowEntry.getKey())
&& (isSameAsPrevious(lowEntry.getKey(), updatedCapacity))) {
cumulativeCapacity.remove(lowEntry.getKey());
} else {
cumulativeCapacity.put(startKey, updatedCapacity);
}
}
// Increase all the capacities of overlapping intervals
Set<Entry<Long, Resource>> overlapSet =
ticks.tailMap(startKey, false).entrySet();
for (Entry<Long, Resource> entry : overlapSet) {
updatedCapacity = Resources.add(entry.getValue(), totCap);
entry.setValue(updatedCapacity);
}
} else {
// This is the first interval to be added
cumulativeCapacity.put(startKey, totCap);
}
Resource nextTick = cumulativeCapacity.get(endKey);
if (nextTick != null) {
// If there is overlap, remove the duplicate entry
if (isSameAsPrevious(endKey, nextTick)) {
cumulativeCapacity.remove(endKey);
}
} else {
// Decrease capacity as this is end of the interval
cumulativeCapacity.put(endKey, Resources.subtract(cumulativeCapacity
.floorEntry(endKey).getValue(), totCap));
}
return true;
} finally {
writeLock.unlock();
}
}
/**
* Removes a resource for the specified interval
*
* @param reservationInterval the interval for which the resource is to be
* removed
* @param totCap the resource to be removed
* @return true if removal is successful, false otherwise
*/
public boolean removeInterval(ReservationInterval reservationInterval,
Resource totCap) {
if (totCap.equals(ZERO_RESOURCE)) {
return true;
}
writeLock.lock();
try {
long startKey = reservationInterval.getStartTime();
long endKey = reservationInterval.getEndTime();
// update the start key
NavigableMap<Long, Resource> ticks =
cumulativeCapacity.headMap(endKey, false);
// Decrease all the capacities of overlapping intervals
SortedMap<Long, Resource> overlapSet = ticks.tailMap(startKey);
if (overlapSet != null && !overlapSet.isEmpty()) {
Resource updatedCapacity = Resource.newInstance(0, 0);
long currentKey = -1;
for (Iterator<Entry<Long, Resource>> overlapEntries =
overlapSet.entrySet().iterator(); overlapEntries.hasNext();) {
Entry<Long, Resource> entry = overlapEntries.next();
currentKey = entry.getKey();
updatedCapacity = Resources.subtract(entry.getValue(), totCap);
// update each entry between start and end key
cumulativeCapacity.put(currentKey, updatedCapacity);
}
// Remove the first overlap entry if it is same as previous after
// updation
Long firstKey = overlapSet.firstKey();
if (isSameAsPrevious(firstKey, overlapSet.get(firstKey))) {
cumulativeCapacity.remove(firstKey);
}
// Remove the next entry if it is same as end entry after updation
if ((currentKey != -1) && (isSameAsNext(currentKey, updatedCapacity))) {
cumulativeCapacity.remove(cumulativeCapacity.higherKey(currentKey));
}
}
return true;
} finally {
writeLock.unlock();
}
}
/**
* Returns the capacity, i.e. total resources allocated at the specified point
* of time
*
* @param tick the time (UTC in ms) at which the capacity is requested
* @return the resources allocated at the specified time
*/
public Resource getCapacityAtTime(long tick) {
readLock.lock();
try {
Entry<Long, Resource> closestStep = cumulativeCapacity.floorEntry(tick);
if (closestStep != null) {
return Resources.clone(closestStep.getValue());
}
return Resources.clone(ZERO_RESOURCE);
} finally {
readLock.unlock();
}
}
/**
* Get the timestamp of the earliest resource allocation
*
* @return the timestamp of the first resource allocation
*/
public long getEarliestStartTime() {
readLock.lock();
try {
if (cumulativeCapacity.isEmpty()) {
return -1;
} else {
return cumulativeCapacity.firstKey();
}
} finally {
readLock.unlock();
}
}
/**
* Get the timestamp of the latest resource allocation
*
* @return the timestamp of the last resource allocation
*/
public long getLatestEndTime() {
readLock.lock();
try {
if (cumulativeCapacity.isEmpty()) {
return -1;
} else {
return cumulativeCapacity.lastKey();
}
} finally {
readLock.unlock();
}
}
/**
* Returns true if there are no non-zero entries
*
* @return true if there are no allocations or false otherwise
*/
public boolean isEmpty() {
readLock.lock();
try {
if (cumulativeCapacity.isEmpty()) {
return true;
}
// Deletion leaves a single zero entry so check for that
if (cumulativeCapacity.size() == 1) {
return cumulativeCapacity.firstEntry().getValue().equals(ZERO_RESOURCE);
}
return false;
} finally {
readLock.unlock();
}
}
@Override
public String toString() {
StringBuilder ret = new StringBuilder();
readLock.lock();
try {
if (cumulativeCapacity.size() > THRESHOLD) {
ret.append("Number of steps: ").append(cumulativeCapacity.size())
.append(" earliest entry: ").append(cumulativeCapacity.firstKey())
.append(" latest entry: ").append(cumulativeCapacity.lastKey());
} else {
for (Map.Entry<Long, Resource> r : cumulativeCapacity.entrySet()) {
ret.append(r.getKey()).append(": ").append(r.getValue())
.append("\n ");
}
}
return ret.toString();
} finally {
readLock.unlock();
}
}
/**
* Returns the JSON string representation of the current resources allocated
* over time
*
* @return the JSON string representation of the current resources allocated
* over time
*/
public String toMemJSONString() {
StringWriter json = new StringWriter();
JsonWriter jsonWriter = new JsonWriter(json);
readLock.lock();
try {
jsonWriter.beginObject();
// jsonWriter.name("timestamp").value("resource");
for (Map.Entry<Long, Resource> r : cumulativeCapacity.entrySet()) {
jsonWriter.name(r.getKey().toString()).value(r.getValue().toString());
}
jsonWriter.endObject();
jsonWriter.close();
return json.toString();
} catch (IOException e) {
// This should not happen
return "";
} finally {
readLock.unlock();
}
}
/**
* Returns the representation of the current resources allocated over time as
* an interval map.
*
* @return the representation of the current resources allocated over time as
* an interval map.
*/
public Map<ReservationInterval, Resource> toIntervalMap() {
readLock.lock();
try {
Map<ReservationInterval, Resource> allocations =
new TreeMap<ReservationInterval, Resource>();
// Empty
if (isEmpty()) {
return allocations;
}
Map.Entry<Long, Resource> lastEntry = null;
for (Map.Entry<Long, Resource> entry : cumulativeCapacity.entrySet()) {
if (lastEntry != null) {
ReservationInterval interval =
new ReservationInterval(lastEntry.getKey(), entry.getKey());
Resource resource = lastEntry.getValue();
allocations.put(interval, resource);
}
lastEntry = entry;
}
return allocations;
} finally {
readLock.unlock();
}
}
}