package com.atlassian.bitbucket.mesh.rpc;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * Helper that registers detached streams in a {@link ThreadLocal} and closes them on demand.
 *
 * @see ZeroCopyServerCallHandler
 * @see ZeroCopyMarshaller
 */
class ZeroCopyStreamCloser {

    public static final ThreadLocal<List<InputStream>> DETACHED_STREAMS = new ThreadLocal<>();

    private static final Logger log = LoggerFactory.getLogger(ZeroCopyStreamCloser.class);

    public static void close() {
        List<InputStream> streams = DETACHED_STREAMS.get();
        try {
            if (streams != null) {
                for (InputStream inputStream : streams) {
                    try {
                        inputStream.close();
                    } catch (Exception e) {
                        log.error("Failed to close detached gRPC input stream. Memory may leak!", e);
                    }
                }
            }
        } finally {
            DETACHED_STREAMS.remove();
        }
    }

    static void register(InputStream inputStream) {
        List<InputStream> streams = DETACHED_STREAMS.get();
        if (streams == null) {
            streams = new ArrayList<>();
            DETACHED_STREAMS.set(streams);
        }
        streams.add(inputStream);
    }
}
