Create R8 specific rules to allow for service loader optimizations

[mkj-gram]

Creating R8 specific rules will allow Proguard to continue to work while allow R8 to remove the ServiceLoaders.

[dkhalanskyjb]

Hi! I didn't follow the issue closely, so please excuse me if I'm asking something obvious, but what does this accomplish? According to #1231 (comment) and the comment below, R8 already knows how to optimize the service loader calls away, even without the proposed fix.

[mkj-gram]

This PR creates r8 specific rules that gradle will use to invoke R8 with instead of the standard ones from:

kotlinx-coroutines-core/jvm/resources/META-INF/proguard/coroutines.pro

The difference between the two files is the presence of the following two rules in proguard/coroutines.pro:

# ServiceLoader support
-keepnames class kotlinx.coroutines.internal.MainDispatcherFactory {}
-keepnames class kotlinx.coroutines.CoroutineExceptionHandler {}

These are added to PG to ensure that the service loaders work - but for R8 it ends up putting a keep rule on them such that the optimization mentioned in #1270 fails to apply after updates to R8. To ensure that the optimization still works in R8 the rules have to be removed. To ensure that PG continue to work as expected I therefore separated the rules out.

[dkhalanskyjb]

Yeah, I understand what this PR does, the question is why.

the optimization mentioned in #1270 fails to apply after updates to R8

Which optimization does that issue mention? Could you link to a specific comment, for example? As far as I can see, it regards errors with which R8 fails, and doesn't at all regard service loaders.

Also, do you mean that the calls to ServiceLoader suddenly stopped being optimized? If so, could you provide a small reproducing project or at least give some specific versions on which the optimization fails to work so I could check it?

[mkj-gram]

Ah sorry, I got confused by the hashtags and numbers. The optimization mentioned in #1231 is an R8 specific optimization that rewrites ServiceLoader.loads on a specific syntactic form to instantiations. In the optimization there is a bailout to prevent R8 from doing that optimization - and that is adding a keep rule to the configuration. If the interface is kept by a rule in any way R8 will not rewrite the serviceloader.

Due to a change in the semantics of -keepnames R8 ends up not rewriting the mentioned service loaders. That is present from the bundled R8 version in AGP 7.0.0. A reproduction was reported at https://issuetracker.google.com/issues/196302685

The -keepnames are not needed anyway because R8 will always regard classes mentioned in a ServiceLoader configuration file as being live if the interface becomes live.

[mkubiczek]

Can we please have this fix integrated and new version of coroutines with this fix released?

[dkhalanskyjb]

I attempted to apply this fix. Looks like it doesn't do the trick?

The bytecode of `MainDispatcherLoader` in the release APK of the reproducing project, using a version of coroutines with the fix applied

.method public static constructor <clinit>()V
    .registers 6

    const-string v0, "kotlinx.coroutines.fast.service.loader"

    .line 1
    invoke-static {v0}, Lkotlin/internal/ProgressionUtilKt;->kotlinx.coroutines.internal.SystemPropsKt.systemProp(Ljava/lang/String;)Ljava/lang/String;

    move-result-object v0

    if-nez v0, :cond_9

    goto :goto_c

    :cond_9
    invoke-static {v0}, Ljava/lang/Boolean;->parseBoolean(Ljava/lang/String;)Z

    .line 2
    :goto_c
    const-class v0, Lkotlinx/coroutines/internal/MainDispatcherFactory;

    .line 3
    :try_start_e
    invoke-virtual {v0}, Ljava/lang/Class;->getClassLoader()Ljava/lang/ClassLoader;

    move-result-object v1

    .line 4
    invoke-static {v0, v1}, Ljava/util/ServiceLoader;->load(Ljava/lang/Class;Ljava/lang/ClassLoader;)Ljava/util/ServiceLoader;

    move-result-object v0

    .line 5
    invoke-virtual {v0}, Ljava/util/ServiceLoader;->iterator()Ljava/util/Iterator;

    move-result-object v0

    invoke-static {v0}, Lkotlin/sequences/SequencesKt___SequencesJvmKt;->kotlin.sequences.SequencesKt__SequencesKt.asSequence(Ljava/util/Iterator;)Lkotlin/sequences/Sequence;

    move-result-object v0

    invoke-static {v0}, Lkotlin/sequences/SequencesKt;->kotlin.sequences.SequencesKt___SequencesKt.toList(Lkotlin/sequences/Sequence;)Ljava/util/List;

    move-result-object v0

    .line 6
    invoke-interface {v0}, Ljava/lang/Iterable;->iterator()Ljava/util/Iterator;

    move-result-object v1

    .line 7
    invoke-interface {v1}, Ljava/util/Iterator;->hasNext()Z

    move-result v2

    if-nez v2, :cond_2e

    const/4 v1, 0x0

    goto :goto_57

    .line 8
    :cond_2e
    invoke-interface {v1}, Ljava/util/Iterator;->next()Ljava/lang/Object;

    move-result-object v2

    .line 9
    invoke-interface {v1}, Ljava/util/Iterator;->hasNext()Z

    move-result v3

    if-nez v3, :cond_3a

    :goto_38
    move-object v1, v2

    goto :goto_57

    .line 10
    :cond_3a
    move-object v3, v2

    check-cast v3, Lkotlinx/coroutines/internal/MainDispatcherFactory;

    .line 11
    invoke-interface {v3}, Lkotlinx/coroutines/internal/MainDispatcherFactory;->getLoadPriority()I

    move-result v3

    .line 12
    :cond_41
    invoke-interface {v1}, Ljava/util/Iterator;->next()Ljava/lang/Object;

    move-result-object v4

    .line 13
    move-object v5, v4

    check-cast v5, Lkotlinx/coroutines/internal/MainDispatcherFactory;

    .line 14
    invoke-interface {v5}, Lkotlinx/coroutines/internal/MainDispatcherFactory;->getLoadPriority()I

    move-result v5

    if-ge v3, v5, :cond_50

    move-object v2, v4

    move v3, v5

    .line 15
    :cond_50
    invoke-interface {v1}, Ljava/util/Iterator;->hasNext()Z

    move-result v4

    if-nez v4, :cond_41

    goto :goto_38

    .line 16
    :goto_57
    check-cast v1, Lkotlinx/coroutines/internal/MainDispatcherFactory;
    :try_end_59
    .catchall {:try_start_e .. :try_end_59} :catchall_6f

    if-eqz v1, :cond_67

    .line 17
    :try_start_5b
    invoke-interface {v1, v0}, Lkotlinx/coroutines/internal/MainDispatcherFactory;->createDispatcher(Ljava/util/List;)Lkotlinx/coroutines/MainCoroutineDispatcher;

    move-result-object v0
    :try_end_5f
    .catchall {:try_start_5b .. :try_end_5f} :catchall_62

    .line 18
    sput-object v0, Lkotlinx/coroutines/internal/MainDispatcherLoader;->dispatcher:Lkotlinx/coroutines/MainCoroutineDispatcher;

    return-void

    :catchall_62
    move-exception v0

    .line 19
    :try_start_63
    invoke-interface {v1}, Lkotlinx/coroutines/internal/MainDispatcherFactory;->hintOnError()Ljava/lang/String;

    .line 20
    throw v0

    .line 21
    :cond_67
    new-instance v0, Ljava/lang/IllegalStateException;

    const-string v1, "Module with the Main dispatcher is missing. Add dependency providing the Main dispatcher, e.g. \'kotlinx-coroutines-android\' and ensure it has the same version as \'kotlinx-coroutines-core\'"

    invoke-direct {v0, v1}, Ljava/lang/IllegalStateException;-><init>(Ljava/lang/String;)V

    throw v0
    :try_end_6f
    .catchall {:try_start_63 .. :try_end_6f} :catchall_6f

    :catchall_6f
    move-exception v0

    .line 22
    throw v0
.end method

[mkj-gram]

You are right, apparently AGP needs com.android.tools/r8 under META-INF to select the correct rules. I've updated the PR.

To test that it actually works I created a test project and attached it here. Under app/libs there is:
jetified-kotlinx-coroutines-core-jvm-1.3.9_original.jar and
jetified-kotlinx-coroutines-core-jvm-1.3.9_r8.jar

where compiling by ./gradlew clean assembleRelease with the r8 dependency successfully removes the ServiceLoader but having original will not.

kotlin-coroutines-test.zip

[dkhalanskyjb]

I can confirm that moving the r8/coroutines.pro to com.android.tools does fix the issue, but now I'm confused about why com.android.tools/proguard/coroutines.pro also needs to be added alongside just proguard/coroutines.pro. Could you please comment on that?

[mkj-gram]

I am just following the same format as here:
https://github.com/Kotlin/kotlinx.coroutines/tree/master/ui/kotlinx-coroutines-android/resources/META-INF/com.android.tools

I would assume it is to pass PG specific rules to PG and R8 specific rules to R8.

In kotlinx.coroutines/ui/kotlinx-coroutines-android/resources/META-INF/proguard/coroutines.pro, it says in the header:

# Files in this directory will be ignored starting with Android Gradle Plugin 3.6.0+

However, I am not an AGP expert.

[nevack]

We had severe performance regressions after migrating to Android Gradle Plugin 7.0.
Version of R8 not changed with this upgrade, because we have already pinned it at 3.0.73.
To fix regression we had to build custom artifact of org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.5.0 with patch from this PR applied.

Are there plans to merge this PR? The problem is very bad and should be fixed ASAP, as AGP 7 is current stable release for at least a month.

[mcomella]

@nevack We were able to work around this issue by using the fast service loader implementation built into the coroutines lib.

tldr; add this to your proguard-rules.pro:

-assumenosideeffects class kotlinx.coroutines.internal.MainDispatcherLoader {
    boolean FAST_SERVICE_LOADER_ENABLED return true;
}

Explanation

The coroutines lib defines two implementations to get the MainDispatcherFactory:

kotlinx.coroutines/kotlinx-coroutines-core/jvm/src/internal/MainDispatchers.kt

Lines 26 to 36 in 8c98180

	val factories = if (FAST_SERVICE_LOADER_ENABLED) {
	FastServiceLoader.loadMainDispatcherFactory()
	} else {
	// We are explicitly using the
	// `ServiceLoader.load(MyClass::class.java, MyClass::class.java.classLoader).iterator()`
	// form of the ServiceLoader call to enable R8 optimization when compiled on Android.
	ServiceLoader.load(
	MainDispatcherFactory::class.java,
	MainDispatcherFactory::class.java.classLoader
	).iterator().asSequence().toList()
	}

By default, it takes the latter branch. This is slow and R8 is supposed to optimize it out. However, in AGP 7.0, this functionality broke: see https://issuetracker.google.com/issues/196302685, which is the bug that caused this issue to be filed.

However, the coroutines library has an alternative implementation – the fast service loader – which is just as fast as the R8 optimized version in our experience. To activate that, you can use proguard to rewrite the boolean that defines the fast service loader here:

kotlinx.coroutines/kotlinx-coroutines-core/jvm/src/internal/MainDispatchers.kt

Line 19 in 8c98180

private val FAST_SERVICE_LOADER_ENABLED = systemProp(FAST_SERVICE_LOADER_PROPERTY_NAME, true)

to true (I couldn't get the system property working). That's the tldr; up top:

-assumenosideeffects class kotlinx.coroutines.internal.MainDispatcherLoader {
    boolean FAST_SERVICE_LOADER_ENABLED return true;
}

Here's the PR where we implemented this change: mozilla-mobile/fenix#20844 I apologize for not documenting this sooner.

[mxalbert1996]

This workaround doesn't prevent SL call for CoroutineExceptionHandler, does it?

[mcomella]

@mxalbert1996 You are right that the workaround I suggested doesn't optimize the ServiceLoader in the CoroutineExceptionHandler – good find. fwiw, I suspect the workaround will be good enough for many users because afaict the CoroutineExceptionHandler SL will only cause a user-perceptible delay the first time an uncaught exception is thrown in a coroutine, which is less likely to happen in a perf critical section than the MainDispatcherFactory SL (which I suspect affects many apps' start up or early UI interactions).

[dkhalanskyjb]

Thanks! Merged this manually in dfc4821

[Thomas-Vos]

@dkhalanskyjb thanks for merging the PR. Could you please publish a bug fix release for 1.5.x which includes this fix?

[dkhalanskyjb]

@Thomas-Vos Kotlin 1.6 will be quite soon, and we plan to publish a new release candidate shortly thereafter, which will include this fix.