AOP w połączeniu z Kotlin Coroutines

Programowanie aspektowe (aspect-oriented programming) jest od dłuższego czasu popularne w świecie javowym. Biblioteki takie jak AspectJ czy Spring AOP pozwalają w łatwy sposób obsłużyć cross-cutting concerns. Niestety, nie wspierają one oficjalnie Kotlin Coroutines, które świetnie sprawdzają się w pisaniu asynchronicznego, nieblokującego kodu. Czy jesteśmy w stanie zmusić aspekty do działania z korutynami, pomimo braku oficjalnego wsparcia? Okazuje się, że tak.

Problemy w łączeniu AOP z asynchronicznym kodem

Tradycyjnie w tych bibliotekach mamy dostępne rady (advices) umożliwiające zdefiniowanie kodu, który ma zostać wykonany na przykład przed wywołaniem metody (before), po zwróceniu wartości przez metodę (after returning) lub po rzuceniu wyjątku (after throwing).

class Foo {
    @Logging
    fun add(a: Int, b: Int): Int = a + b
}

@Aspect
class LoggingAspect {
    @AfterReturning(
        pointcut = "@annotation(Logging)",
        returning = "result"
    )
    fun logResult(joinPoint: JoinPoint, result: Any?) {
        logger.debug("`${joinPoint.signature}` returned `${result}`")
    }
}

[DEBUG] `int Foo.add(int,int)` returned `5`

O ile w tradycyjnym modelu równoległości opartym o wątki i blokujące wywołania sprawdza się to dobrze, o tyle w nowszych podejściach opartych o asynchroniczność jest już gorzej. Jeśli używamy Reactive Streams, często będziemy mieli metody, które zamiast bezpośrednio zwracać wartości i rzucać wyjątki, zwracają leniwe strumienie Publisher<T>, które to emitują wartości lub błędy (onNext/onError). Pisanie aspektów staje się trochę uciążliwe, lecz wciąż wykonalne:

class ReactiveFoo {
    @ReactiveLogging
    fun add(a: Int, b: Int): Mono<Int> = Mono.just(a + b)
}

@Aspect
class ReactiveLoggingAspect {
    @Around("""
        @annotation(ReactiveLogging) &&
        execution(reactor.core.publisher.Mono *(..))
    """)
    fun logResult(joinPoint: ProceedingJoinPoint): Mono<Any> =
        (joinPoint.proceed() as Mono<Any>)
            .doOnNext { result -> 
                logger.debug(
                    "`${joinPoint.signature}` returned `${result}`"
                )
            }
}

[DEBUG] `Mono ReactiveFoo.add(int,int)` returned `5`

Gorzej jest w przypadku Kotlin Coroutines. AspectJ/Spring AOP oficjalnie nie wspierają suspending functions. Na pierwszy rzut oka asynchroniczny kod oparty o korutyny niewiele różni się od klasycznego kodu synchronicznego. Widzimy słowo kluczowe suspend i niewiele poza tym. Kompilator Kotlina jednak dość mocno przekształca ten kod i na poziomie bytecodu otrzymujemy zmienione sygnatury metod, które wykorzystują Continuation Passing Style i implementują maszyny stanów.

Spróbujmy zaaplikować standardowy aspekt do suspend fun:

class CoroutineFoo {
    @Logging
    suspend fun add(a: Int, b: Int): Int {
        delay(100)
        return a + b
    }
}

Po uruchomieniu sygnatura funkcji, jak i zwracana wartość mogą się wydawać inne od oczekiwanych:

[DEBUG] `Object CoroutineFoo.add(int,int,Continuation)` returned `COROUTINE_SUSPENDED`

Każda suspend funkcja na poziomie bytecodu JVM ma dodatkowy parametr Continuation. Jej zwracany typ to Object, ponieważ oprócz typu zadeklarowanego w Kotlinie, może też zwrócić magiczną wartość COROUTINE_SUSPENDED.

Jak zmusić aspekty do działania z korutynami?

Zauważmy, że sygnatury funkcji suspend fun f(): T oraz fun f(c: Continuation): Any? z punktu widzenia JVM są identyczne. Z punktu widzenia Kotlina, w pierwszej Continuation jest przekazywane implicite, w drugiej explicite. Pierwsza nie jest wspierana przez AspectJ, ale druga już jest.

Do zdefiniowania aspektu wykorzystamy punkt cięcia (pointcut) args(.., kotlin.coroutines.Continuation) oraz radę (advice) around:

@Around("""
    @annotation(CoroutineLogging) &&
    args(.., kotlin.coroutines.Continuation)
""")
fun logResult(joinPoint: ProceedingJoinPoint): Any? { }

Z ProceedingJoinPoint potrzebujemy wyciągnąć argumenty, z którymi została wywołana funkcja, w szczególności ostatni argument typu Continuation:

val ProceedingJoinPoint.coroutineContinuation: Continuation<Any?>  
    get() = this.args.last() as Continuation<Any?>  

val ProceedingJoinPoint.coroutineArgs: Array<Any?>  
    get() = this.args.sliceArray(0 until this.args.size - 1)

Pozwala nam to na zaimplementowanie odpowiednika ProceedingJoinPoint.proceed() zdefiniowanego jako suspend fun:

import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn

suspend fun ProceedingJoinPoint.proceedCoroutine(
    args: Array<Any?> = this.coroutineArgs
): Any? =
    suspendCoroutineUninterceptedOrReturn { continuation ->  
        this.proceed(args + continuation)  
    }

Wykorzystując ProceedingJoinPoint.coroutineContinuation, które zdefiniowaliśmy chwilę wcześniej, możemy wywołać dowolną funkcję suspend:

import kotlin.coroutines.intrinsics.startCoroutineUninterceptedOrReturn

fun ProceedingJoinPoint.runCoroutine(
    block: suspend () -> Any?
): Any? =
    block.startCoroutineUninterceptedOrReturn(this.coroutineContinuation)

Składając to wszystko razem, otrzymujemy aspekt działający z korutynami:

@Aspect
class CoroutineLoggingAspect {  
    @Around("""
        @annotation(CoroutineLogging) &&
        args(.., kotlin.coroutines.Continuation)
    """)
    fun logResult(joinPoint: ProceedingJoinPoint): Any? =  
        joinPoint.runCoroutine {
            val result = joinPoint.proceedCoroutine()  
            logger.debug("`${joinPoint.signature}` returned `${result}`")  
            result
        }
}

[DEBUG] `Object CoroutineFoo.add(int,int,Continuation)` returned `5`

Pełny kod źródłowy znajdziecie tutaj.

Podsumowanie

Używanie aspektów z korutynami jest jak najbardziej możliwe. Pomimo braku oficjalnego wsparcia przy odrobinie dodatkowego kodu jesteśmy w stanie połączyć programowanie aspektowe z modelem asynchronicznym. Aspekty mogą nam się przydać przy implementowaniu cross-cutting concerns takich jak logging, tracing i tym podobne.

Problemy w łączeniu AOP z asynchronicznym kodem

Jak zmusić aspekty do działania z korutynami?

Podsumowanie

Odkryj więcej świetnych treści!