org.apache.flink.runtime.memory.SharedResources Maven / Gradle / Ivy

Go to download
/*
 * 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.flink.runtime.memory;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.util.function.LongFunctionWithException;

import javax.annotation.concurrent.GuardedBy;

import java.util.HashMap;
import java.util.HashSet;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.LongConsumer;

import static org.apache.flink.util.Preconditions.checkState;

/** A map that keeps track of acquired shared resources and handles their allocation disposal. */
public final class SharedResources {

    private final ReentrantLock lock = new ReentrantLock();

    @GuardedBy("lock")
    private final HashMap> reservedResources = new HashMap<>();

    /**
     * Gets the shared memory resource for the given owner and registers a lease. If the resource
     * does not yet exist, it will be created via the given initializer function.
     *
     * The resource must be released when no longer used. That releases the lease. When all
     * leases are released, the resource is disposed.
     */
    public  ResourceAndSize getOrAllocateSharedResource(
            String type,
            Object leaseHolder,
            LongFunctionWithException initializer,
            long sizeForInitialization)
            throws Exception {

        // We could be stuck on this lock for a while, in cases where another initialization is
        // currently
        // happening and the initialization is expensive.
        // We lock interruptibly here to allow for faster exit in case of cancellation errors.
        try {
            lock.lockInterruptibly();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new MemoryAllocationException("Interrupted while acquiring memory");
        }

        try {
            // we cannot use "computeIfAbsent()" here because the computing function may throw an
            // exception.
            @SuppressWarnings("unchecked")
            LeasedResource resource = (LeasedResource) reservedResources.get(type);
            if (resource == null) {
                resource = createResource(initializer, sizeForInitialization);
                reservedResources.put(type, resource);
            }

            resource.addLeaseHolder(leaseHolder);
            return resource;
        } finally {
            lock.unlock();
        }
    }

    /**
     * Releases a lease (identified by the lease holder object) for the given type. If no further
     * leases exist, the resource is disposed.
     */
    void release(String type, Object leaseHolder) throws Exception {
        release(type, leaseHolder, (value) -> {});
    }

    /**
     * Releases a lease (identified by the lease holder object) for the given type. If no further
     * leases exist, the resource is disposed.
     *
     * 
This method takes an additional hook that is called when the resource is disposed.
     */
    public void release(String type, Object leaseHolder, LongConsumer releaser) throws Exception {
        lock.lock();
        try {
            final LeasedResource resource = reservedResources.get(type);
            if (resource == null) {
                return;
            }

            if (resource.removeLeaseHolder(leaseHolder)) {
                try {
                    reservedResources.remove(type);
                    resource.dispose();
                } finally {
                    releaser.accept(resource.size());
                }
            }
        } finally {
            lock.unlock();
        }
    }

    @VisibleForTesting
    int getNumResources() {
        return reservedResources.size();
    }

    private static  LeasedResource createResource(
            LongFunctionWithException initializer, long size) throws Exception {

        final T resource = initializer.apply(size);
        return new LeasedResource<>(resource, size);
    }

    // ------------------------------------------------------------------------

    /** A resource handle with size. */
    public interface ResourceAndSize {

        T resourceHandle();

        long size();
    }

    // ------------------------------------------------------------------------

    private static final class LeasedResource
            implements ResourceAndSize {

        private final HashSet