java-shade.gradle

import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
import proguard.gradle.ProGuardTask

import java.util.concurrent.atomic.AtomicInteger

buildscript {
    repositories {
        gradlePluginPortal()
        google()
    }
    dependencies {
        classpath "com.gradleup.shadow:com.gradleup.shadow.gradle.plugin:${managedVersions['com.gradleup.shadow:com.gradleup.shadow.gradle.plugin']}"
        classpath "com.guardsquare:proguard-gradle:${managedVersions['com.guardsquare:proguard-gradle']}"
    }
}

def globalShadowExclusions = []
if (rootProject.hasProperty('shadowExclusions')) {
    def rawShadowExclusions = rootProject.findProperty("shadowExclusions")
    if (!(rawShadowExclusions instanceof String)) {
        throw new RuntimeException("Property 'shadowExclusions' must be a String but was: ${rawShadowExclusions.getClass().name}")
    }
    globalShadowExclusions = rawShadowExclusions.split(",")
}
subprojects {
    ext {
        shadowExclusions = globalShadowExclusions.collect()
    }
}

def relocatedProjects = projectsWithFlags('java', 'relocate')
def numConfiguredRelocatedProjects = new AtomicInteger()
configure(relocatedProjects) {

    // per-dependency shade configuration
    // while shade per-project shades all dependencies, shade per-dependency allows for finer control
    def shadeConfig = project.configurations.create("shade") {
        canBeResolved = true
        canBeConsumed = false
    }
    project.configurations.api.extendsFrom(shadeConfig)

    // Generate the shaded JARs.
    task shadedJar(
            type: ShadowJar,
            group: 'Build',
            description: 'Builds the shaded main JAR.',
            dependsOn: tasks.classes) {

        def self = this
        def delegate0 = delegate
        project.afterEvaluate {
            self.configureShadowTask(project, delegate0, true)
        }
        archiveClassifier.set('shaded')

        // Exclude the legacy file listing.
        exclude '/META-INF/INDEX.LIST'
        // Exclude the class signature files.
        exclude '/META-INF/*.SF'
        exclude '/META-INF/*.DSA'
        exclude '/META-INF/*.RSA'
        // Exclude the files generated by Maven
        exclude '/META-INF/maven/**'
        // Exclude the module metadata that'll become invalid after relocation.
        exclude '**/module-info.class'

        // Exclude the MRJARs files that are not compatible with the target Java version.
        exclude { details ->
            def path = details.path  // e.g., META-INF/versions/15/com/example/Foo.class
            def matcher = path =~ /^META-INF\/versions\/(\d+)\/.*$/
            if (matcher.matches()) {
                def version = matcher[0][1].toInteger()
                return version > project.ext.targetJavaVersion
            }
            return false
        }

        // Set the 'Automatic-Module-Name' property in MANIFEST.MF.
        if (project.ext.automaticModuleName != null) {
            doFirst {
                manifest {
                    attributes('Automatic-Module-Name': project.ext.automaticModuleName.get())
                }
            }
        }
    }
    tasks.assemble.dependsOn tasks.shadedJar

    artifacts {
        archives shadedJar
    }

    task shadedClasses(
            type: Copy,
            group: 'Build',
            description: 'Extracts the shaded main JAR.',
            dependsOn: tasks.shadedJar) {

        from(zipTree(tasks.shadedJar.archiveFile.get().asFile))
        from(sourceSets.main.output.classesDirs) {
            // Add the JAR resources excluded in the 'shadedJar' task.
            include '**/*.jar'
        }
        into "${project.buildDir}/classes/java/shaded-main"
    }

    task shadedTestJar(
            type: ShadowJar,
            group: 'Build',
            description: 'Builds the shaded test JAR.',
            dependsOn: tasks.testClasses) {

        def self = this
        def delegate0 = delegate
        project.afterEvaluate {
            self.configureShadowTask(project, delegate0, false)
        }
        archiveBaseName.set("test-${tasks.jar.archiveBaseName.get()}-shaded")
    }

    task copyShadedTestClasses(
            type: Copy,
            group: 'Build',
            description: 'Extracts the shaded test JAR.',
            dependsOn: tasks.shadedTestJar) {
        from(zipTree(tasks.shadedTestJar.archiveFile.get().asFile))
        from(sourceSets.test.output.classesDirs) {
            // Add the JAR resources excluded in the 'shadedTestJar' task.
            include '**/*.jar'
        }
        into "${project.buildDir}/classes/java/shaded-test"
    }
}

// NB: Configure in a new closure so that all relocated projects have a 'shadedJar' task.
configure(relocatedProjects) {
    if (project.hasFlags('trim')) {
        // Task 'shadedJar' may produce a very large JAR. Rename it to '<artifact>-untrimmed-<version>-shaded.jar' and
        // let the task 'trimShadedJar' produce the trimmed JAR from it.
        tasks.shadedJar.archiveBaseName.set("${tasks.jar.archiveBaseName.get()}-untrimmed")

        task trimShadedJar(
                type: ProGuardTask,
                group: 'Build',
                description: 'Shrinks the shaded JAR by removing unused classes.') {

            relocatedProjects.each {
                dependsOn it.tasks.shadedJar.path
                dependsOn it.tasks.shadedTestJar.path
            }

            def shadedFile = tasks.shadedJar.archiveFile.get().asFile
            // Rename the output file to '<artifact>-<version>-shaded.jar' by removing '-untrimmed' from
            // the input file name.
            def shadedAndTrimmedFile = file(shadedFile.path.replaceFirst('-untrimmed-', '-'))

            injars shadedFile
            // NB: By specifying 'outjars' *before* other 'injars' below, ProGuard will put only the classes
            //     from 'shadedFile' into 'shadedAndTrimmedFile'. See 'restructuring the output archives'
            //     for more information: https://www.guardsquare.com/en/proguard/manual/examples#restructuring
            outjars shadedAndTrimmedFile

            // Include all other shaded JARs so that ProGuard does not trim the classes and methods
            // that are used actually.
            injars tasks.shadedTestJar.archiveFile.get().asFile
            relocatedProjects.each {
                if (it != project && !it.hasFlags('no_aggregation')) {
                    injars it.tasks.shadedJar.archiveFile.get().asFile
                    injars it.tasks.shadedTestJar.archiveFile.get().asFile
                }
            }

            def proguardLibraryConfigs = []
            relocatedProjects.each { p ->

                if (!p.hasFlags('no_aggregation')) {
                    configure(p) {
                        configurations {
                            proguardLibraryJarsElements {
                                canBeConsumed = true
                                canBeResolved = false
                                extendsFrom configurations.dependencyManagement
                                extendsFrom configurations.runtimeClasspath
                                attributes {
                                    attribute(Usage.USAGE_ATTRIBUTE, objects.named(Usage, Usage.JAVA_RUNTIME))
                                }
                            }
                        }
                    }

                    def configName = "proguardLibraryJars_${p.path.replace(':', '_')}"
                    def proguardLibraryJars = configurations.create(configName) {
                        canBeConsumed = false
                        canBeResolved = true
                        // add attributes to match producer if you set any
                        attributes {
                            attribute(Usage.USAGE_ATTRIBUTE, objects.named(Usage, Usage.JAVA_RUNTIME))
                        }
                    }

                    proguardLibraryConfigs.add(proguardLibraryJars)
                    dependencies {
                        add(proguardLibraryJars.getName(), project(path: p.path, configuration: 'proguardLibraryJarsElements'))
                    }
                }
            }

            // Filtering the libraryjars needs to be deferred to ensure all configurations are fully populated.
            doFirst {
                def dependencyJars = new LinkedHashSet<File>()
                proguardLibraryConfigs.each { config ->
                    config.each { File file ->
                        if (!file.path.startsWith("${rootProject.projectDir}")) {
                            dependencyJars.add(file)
                        }
                    }
                }
                libraryjars files(dependencyJars)
            }

            // Add JDK. Use modules instead of rt.jar if possible.
            def javaToolchainService = project.extensions.getByType(JavaToolchainService)
            def javaPluginExt = project.extensions.getByType(JavaPluginExtension)
            def launcher = javaToolchainService.launcherFor(javaPluginExt.toolchain).get()
            File jdkHome = launcher.metadata.installationPath.asFile
            File jmodDir = new File(jdkHome, "jmods")
            if (jmodDir.isDirectory()) {
                jmodDir.listFiles().findAll { File f ->
                    f.isFile() && f.name.toLowerCase(Locale.ENGLISH).endsWith(".jmod")
                }.sort().each { libraryjars it }
            } else {
                File rtJar = new File(jdkHome, "jre/lib/rt.jar")
                libraryjars rtJar
            }

            dontoptimize
            dontobfuscate
            dontwarn // Ignore the harmless 'missing classes' warnings related with the optional dependencies.

            keepattributes '**'
            keepparameternames

            printconfiguration file("${project.buildDir}/proguard.cfg")
        }

        tasks.shadedClasses.dependsOn tasks.trimShadedJar
        tasks.assemble.dependsOn tasks.trimShadedJar

        // Add the trimmed JAR to archives.
        artifacts {
            trimShadedJar.outJarFiles.each {
                archives it
            }
        }
    }
}

// NB: Configure in a new closure so that all relocated projects have a 'shadedJar' or 'trimShadedJar' task.
configure(relocatedProjects) {
    // Make sure the main classes are ready as well when `copyShadedTestClasses` is run,
    // so that the tasks that require `copyShadedTestClasses` have everything required to run shaded tests.
    tasks.copyShadedTestClasses.configure {
        relocatedProjects.each {
            if (it.tasks.findByName('trimShadedJar')) {
                dependsOn it.tasks.trimShadedJar.path
            } else {
                dependsOn it.tasks.shadedJar.path
            }
        }
    }

    // Add tests for the shaded JAR.
    testing.suites {
        shadedTest(JvmTestSuite) {

            targets.configureEach {
                testTask.configure {

                    group = 'Verification'
                    description = 'Runs the unit tests with the shaded classes.'

                    project.ext.configureCommonTestSettings(it)

                    dependsOn tasks.copyShadedTestClasses

                    // The tests against the shaded artifacts should run after the tests against the unshaded ones.
                    shouldRunAfter tasks.test

                    testClassesDirs = files(tasks.copyShadedTestClasses.destinationDir)
                    classpath = testClassesDirs

                    project.ext.relocations.each {
                        exclude "${it['to'].replace('.', '/')}/**"
                    }
                }
            }
        }
    }
    tasks.check.dependsOn testing.suites.shadedTest

    // Configurations shouldn't be created in afterEvaluate so create it eagerly here before configuring it.

    // In Gradle 8, dependencies can't be declared in a resolvable (canBeResolved=true) configuration. As a
    // workaround two configurations are created to separate the declaration and resolving roles.
    // - Directly declares dependencies to `shadedJarTestImplementation`.
    // - Indirectly resolves the dependencies of `shadedJarTestImplementation` by creating a new configuration
    //   and extends `shadedJarTestImplementation`.
    // https://discuss.gradle.org/t/problem-with-compileclasspath-after-gradle-8-update/44940
    project.configurations.create("shadedJarTestImplementation") {
        canBeResolved = false
        canBeConsumed = false
    }
    project.configurations.create("shadedJarTestRuntime") {
        canBeResolved = true
        canBeConsumed = false
        extendsFrom(project.configurations.shadedJarTestImplementation)
        extendsFrom(project.configurations.dependencyManagement)
    }

    // Update the classpath of the 'shadedTest' task after all shaded projects are evaluated
    // so that we get the complete dependency list.
    project.afterEvaluate {
        if (numConfiguredRelocatedProjects.incrementAndGet() == relocatedProjects.size()) {
            relocatedProjects.each { p ->
                configureShadedTestImplementConfiguration(p)
            }
            relocatedProjects.each { p ->
                p.testing.suites.shadedTest.targets.configureEach {
                    testTask.configure {
                        classpath += p.configurations.getByName('shadedJarTestRuntime')
                    }
                }
            }
        }
    }

    gradle.taskGraph.whenReady {
        // Skip unshaded tests if shaded tests will run.
        // To enable, set the property 'preferShadedTests' to 'false'.
        boolean runUnshadedTests = false
        if (rootProject.hasProperty('preferShadedTests') && "false" == rootProject.property('preferShadedTests')) {
            runUnshadedTests = true
        }
        if (gradle.taskGraph.hasTask(tasks.shadedTest)) {
            tasks.test.onlyIf { runUnshadedTests }
        }
    }
}

private void configureShadowTask(Project project, ShadowJar task, boolean isMain) {
    List<Map<String, String>> relocations = project.ext.relocations

    task.configure {
        from(project.sourceSets[isMain ? 'main' : 'test'].output) {
            exclude 'META-INF/maven/**'
            // Prevent the shadow plugin from exploding the JARs in the resources directory.
            // e.g. WEB-INF/lib/hello.jar
            exclude '**/*.jar'
        }
        project.ext.shadowExclusions.each {
            exclude it
        }

        configurations = [project.configurations[isMain ? 'compileClasspath' : 'testCompileClasspath']]

        def shadeConfig = project.configurations.findByName("shade")
        def shadedDeps = shadeConfig ?
                shadeConfig.dependencies.collect { "${it.group}:${it.name}" }.toSet() :
                [] as Set<String>


        relocations.each { props ->
            logger.debug("[${project.name}] Relocating ${props['name']} from ${props['from']} to ${props['to']} " +
                    "in ${isMain ? 'main' : 'test'}")
            task.relocate(props['from'], props['to']) {
                exclude {
                    return isGenerated(project, it.path, isMain)
                }
            }
        }

        dependencies {
            if (project.hasFlags('shade')) {
                // Shade the relocated dependencies only.
                exclude { dep ->
                    if (!relocations.find { dep.name.startsWith("${it['name']}:") }) {
                        // Do not shade the dependencies not listed in 'relocations'.
                        return true
                    }

                    if (isMain) {
                        return false
                    } else {
                        // Do not shade the dependencies which is already shaded for 'main'.
                        return project.configurations.compileClasspath.allDependencies.any { compileDep ->
                            compileDep.group == dep.moduleGroup && compileDep.name == dep.moduleName
                        }
                    }
                }
            } else { // hasFlags('relocate')
                // Only want to shade if explicitly specified on the dependency
                exclude {
                    return !shadedDeps.contains("${it.moduleGroup}:${it.moduleName}")
                }
            }
        }
    }
}

boolean isGenerated(Project project, String path, boolean isMain) {
    if (path.endsWith("package-info.class")) {
        return false
    }
    if (!path.endsWith('.class')) {
        return false
    }
    // Remove the .class extension.
    def className = path.substring(0, path.length() - 6)
    // Remove the inner class part.
    def topLevelClassName = className.split('\\$')[0]

    def javaSourcePath = "${topLevelClassName}.java"
    def scope = isMain ? 'main' : 'test'
    def thriftGenSrcDir = "${project.ext.genSrcDir}/${scope}/javaThrift"
    def grpcGenSrcDir = "${project.buildDir}/generated/sources/proto/${scope}/grpc"
    def protoGenSrcDir = "${project.buildDir}/generated/sources/proto/${scope}/java"

    // Check if the source file exists in any of the generated source directories.
    if (new File("${thriftGenSrcDir}/${javaSourcePath}").exists() ||
        new File("${grpcGenSrcDir}/${javaSourcePath}").exists() ||
        new File("${protoGenSrcDir}/${javaSourcePath}").exists()) {
        logger.debug("[${project.name}] Excluding generated class from shading: ${path}")
        return true
    }
    return false
}

/**
 * Finds the dependencies of {@code project} recursively and adds the found dependencies to
 * the configuration named as {@code 'shadedJarTestImplementation'}.
 */
private Configuration configureShadedTestImplementConfiguration(
        Project project, Project recursedProject = project,
        Set<ExcludeRule> excludeRules = new HashSet<>(),
        Set<Project> visitedProjects = new HashSet<>(),
        boolean recursedProjectRelocated = true) {

    def shadedJarTestImplementation = project.configurations.getByName('shadedJarTestImplementation')

    if (visitedProjects.contains(recursedProject)) {
        return shadedJarTestImplementation
    } else {
        visitedProjects.add(recursedProject);
    }

    if (recursedProject.tasks.findByName('trimShadedJar')) {
        project.dependencies.add(shadedJarTestImplementation.name, files(recursedProject.tasks.trimShadedJar.outJarFiles))
    } else if (recursedProject.tasks.findByName('shadedJar')) {
        project.dependencies.add(shadedJarTestImplementation.name, files(recursedProject.tasks.shadedJar.archiveFile.get().asFile))
    }

    def shadedDependencyNames = project.ext.relocations.collect { it['name'] }
    def targetConfigurationName
    if (project == recursedProject) {
        targetConfigurationName = 'testRuntimeClasspath'
    } else {
        targetConfigurationName = 'runtimeClasspath'
    }
    def projectDependencies = []
    recursedProject.configurations.findAll { cfg ->
        cfg.name == targetConfigurationName
    }.each { cfg ->
        cfg.allDependencies.each { dep ->
            if (dep instanceof ProjectDependency) {
                if (!rootProject.findProject(dep.path).ext.hasFlag('java')) {
                    // Do not add the dependencies of non-Java projects.
                    return
                }
                // Project dependency - recurse later.
                // Note that we recurse later to have immediate module dependencies higher precedence.
                projectDependencies.add(dep)
            } else {
                // Module dependency - add.
                if (shadedDependencyNames.contains("${dep.group}:${dep.name}")) {
                    if (recursedProjectRelocated) {
                        // Skip the shaded dependencies.
                        return
                    }
                    throw new IllegalStateException(
                            "${recursedProject} has a shaded dependency: ${dep.group}:${dep.name} " +
                                    "but it is not relocated. Please add a 'relocate' flag to " +
                                    "${recursedProject} in settings.gradle.")
                }

                if (excludeRules.find { rule ->
                    return rule.group == dep.group && rule.module == dep.name
                }) {
                    // Skip the excluded dependencies.
                    return
                }
                // Do not use `project.dependencies.add(name, dep)` that discards the classifier of
                // a dependency. See https://github.com/gradle/gradle/issues/23096
                shadedJarTestImplementation.dependencies.add(dep)
            }
        }
    }

    // Recurse into the project dependencies.
    projectDependencies.each { ProjectDependency dep ->
        def dependencyProject = rootProject.findProject(dep.path)
        if (!dependencyProject.ext.hasFlag('relocate')) {
            shadedJarTestImplementation.dependencies.add(
                    project.dependencies.project(path: dep.dependencyProject.path))
            recursedProjectRelocated = false
        } else {
            recursedProjectRelocated = true
        }
        configureShadedTestImplementConfiguration(
                project, dependencyProject,
                excludeRules + dep.excludeRules, visitedProjects, recursedProjectRelocated)
    }

    return shadedJarTestImplementation
}