org.testcontainers.images.LoggedPullImageResultCallback Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of testcontainers Show documentation
Isolated container management for Java code testing
There is a newer version: 1.9.1
package org.testcontainers.images;

import com.github.dockerjava.api.command.PullImageResultCallback;
import com.github.dockerjava.api.model.PullResponseItem;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;

import java.io.Closeable;
import java.time.Duration;
import java.time.Instant;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;

/**
 * {@link PullImageResultCallback} with improved logging of pull progress.
 */
class LoggedPullImageResultCallback extends PullImageResultCallback {

    private final Logger logger;

    private final Set allLayers = new HashSet<>();

    private final Set downloadedLayers = new HashSet<>();

    private final Set pulledLayers = new HashSet<>();

    private final Map totalSizes = new HashMap<>();

    private final Map currentSizes = new HashMap<>();

    private boolean completed;

    private Instant start;

    LoggedPullImageResultCallback(final Logger logger) {
        this.logger = logger;
    }

    @Override
    public void onStart(final Closeable stream) {
        super.onStart(stream);
        start = Instant.now();

        logger.info("Starting to pull image");
    }

    @Override
    public void onNext(final PullResponseItem item) {
        super.onNext(item);

        final String statusLowercase = item.getStatus() != null ? item.getStatus().toLowerCase() : "";
        final String id = item.getId();

        if (item.getProgressDetail() != null) {
            allLayers.add(id);
        }

        if (statusLowercase.equalsIgnoreCase("download complete")) {
            downloadedLayers.add(id);
        }

        if (statusLowercase.equalsIgnoreCase("pull complete")) {
            pulledLayers.add(id);
        }

        if (item.getProgressDetail() != null) {
            Long total = item.getProgressDetail().getTotal();
            Long current = item.getProgressDetail().getCurrent();

            if (total != null && total > totalSizes.getOrDefault(id, 0L)) {
                totalSizes.put(id, total);
            }
            if (current != null && current > currentSizes.getOrDefault(id, 0L)) {
                currentSizes.put(id, current);
            }
        }

        if (statusLowercase.startsWith("pulling from") || statusLowercase.contains("complete")) {
            long totalSize = totalLayerSize();
            long currentSize = downloadedLayerSize();

            int pendingCount = allLayers.size() - downloadedLayers.size();
            String friendlyTotalSize;
            if (pendingCount > 0) {
                friendlyTotalSize = "? MB";
            } else {
                friendlyTotalSize = FileUtils.byteCountToDisplaySize(totalSize);
            }

            logger.info(
                "Pulling image layers: {} pending, {} downloaded, {} extracted, ({}/{})",
                String.format("%2d", pendingCount),
                String.format("%2d", downloadedLayers.size()),
                String.format("%2d", pulledLayers.size()),
                FileUtils.byteCountToDisplaySize(currentSize),
                friendlyTotalSize
            );
        }

        if (statusLowercase.contains("complete")) {
            completed = true;
        }
    }

    @Override
    public void onComplete() {
        super.onComplete();

        final long downloadedLayerSize = downloadedLayerSize();
        final long duration = Duration.between(start, Instant.now()).getSeconds();

        if (completed) {
            logger.info(
                "Pull complete. {} layers, pulled in {}s (downloaded {} at {}/s)",
                allLayers.size(),
                duration,
                FileUtils.byteCountToDisplaySize(downloadedLayerSize),
                FileUtils.byteCountToDisplaySize(downloadedLayerSize / duration)
            );
        }
    }

    private long downloadedLayerSize() {
        return currentSizes.values().stream().filter(Objects::nonNull).mapToLong(it -> it).sum();
    }

    private long totalLayerSize() {
        return totalSizes.values().stream().filter(Objects::nonNull).mapToLong(it -> it).sum();
    }
}