Skip to content

Commit 17958c2

Browse files
zach-klippensteinzach-klippenstein
committed
Simplify compose workflow interception, implement in all built-in interceptors.
1 parent 78662ad commit 17958c2

File tree

6 files changed

+166
-112
lines changed

6 files changed

+166
-112
lines changed

‎benchmarks/performance-poetry/complex-poetry/src/main/java/com/squareup/benchmarks/performance/complex/poetry/instrumentation/PerformanceTracingInterceptor.kt

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package com.squareup.benchmarks.performance.complex.poetry.instrumentation
22

3+
import androidx.compose.runtime.Composable
34
import androidx.tracing.Trace
45
import com.squareup.benchmarks.performance.complex.poetry.PerformancePoemWorkflow
56
import com.squareup.benchmarks.performance.complex.poetry.PerformancePoemsBrowserWorkflow
7+
import com.squareup.benchmarks.performance.complex.poetry.instrumentation.PerformanceTracingInterceptor.Companion.NODES_TO_TRACE
68
import com.squareup.workflow1.BaseRenderContext
9+
import com.squareup.workflow1.WorkflowExperimentalApi
710
import com.squareup.workflow1.WorkflowInterceptor
811
import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
912
import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
@@ -27,6 +30,24 @@ class PerformanceTracingInterceptor(
2730
context: BaseRenderContext<P, S, O>,
2831
proceed: (P, S, RenderContextInterceptor<P, S, O>?) -> R,
2932
session: WorkflowSession
33+
): R = traceRender(session) {
34+
proceed(renderProps, renderState, null)
35+
}
36+
37+
@OptIn(WorkflowExperimentalApi::class)
38+
@Composable
39+
override fun <P, O, R> onRenderComposeWorkflow(
40+
renderProps: P,
41+
emitOutput: (O) -> Unit,
42+
proceed: @Composable (P, (O) -> Unit) -> R,
43+
session: WorkflowSession
44+
): R = traceRender(session) {
45+
proceed(renderProps, emitOutput)
46+
}
47+
48+
private inline fun <R> traceRender(
49+
session: WorkflowSession,
50+
render: () -> R
3051
): R {
3152
val isRoot = session.parent == null
3253
val traceIdIndex = NODES_TO_TRACE.indexOfFirst { it.second == session.identifier }
@@ -45,7 +66,7 @@ class PerformanceTracingInterceptor(
4566
Trace.beginSection(sectionName)
4667
}
4768

48-
return proceed(renderProps, renderState, null).also {
69+
return render().also {
4970
if (traceIdIndex > -1 && !sample) {
5071
Trace.endSection()
5172
}

‎workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/SimpleLoggingWorkflowInterceptor.kt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.squareup.workflow1
22

3+
import androidx.compose.runtime.Composable
34
import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
45
import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
56
import kotlinx.coroutines.CoroutineScope
@@ -49,6 +50,17 @@ public open class SimpleLoggingWorkflowInterceptor : WorkflowInterceptor {
4950
proceed(renderProps, renderState, SimpleLoggingContextInterceptor(session))
5051
}
5152

53+
@OptIn(WorkflowExperimentalApi::class)
54+
@Composable
55+
override fun <P, O, R> onRenderComposeWorkflow(
56+
renderProps: P,
57+
emitOutput: (O) -> Unit,
58+
proceed: @Composable (P, (O) -> Unit) -> R,
59+
session: WorkflowSession
60+
): R = logMethod("onRenderComposeWorkflow", session) {
61+
proceed(renderProps, emitOutput)
62+
}
63+
5264
override fun <S> onSnapshotState(
5365
state: S,
5466
proceed: (S) -> Snapshot?,

‎workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/WorkflowInterceptor.kt

Lines changed: 35 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package com.squareup.workflow1
22

3+
import androidx.compose.runtime.Composable
34
import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
45
import com.squareup.workflow1.WorkflowInterceptor.RuntimeLoopOutcome
56
import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
7+
import com.squareup.workflow1.compose.WorkflowComposable
68
import kotlinx.coroutines.CoroutineScope
79
import kotlinx.coroutines.Job
810
import kotlin.coroutines.CoroutineContext
@@ -111,20 +113,6 @@ public interface WorkflowInterceptor {
111113
): RenderingAndSnapshot<R> =
112114
proceed(renderProps)
113115

114-
/**
115-
* Called before a workflow is rendered, and before the interceptor's [onRender] method. Every
116-
* call to this method is eventually followed by a call to [onAfterRender].
117-
*
118-
* This method should be preferred over [onRender] since it supports additional types of workflows
119-
* that will not call [onRender].
120-
*/
121-
public fun <P, S> onBeforeRender(
122-
renderProps: P,
123-
renderState: S,
124-
session: WorkflowSession
125-
) {
126-
}
127-
128116
/**
129117
* Intercepts calls to [StatefulWorkflow.render].
130118
*/
@@ -136,14 +124,15 @@ public interface WorkflowInterceptor {
136124
session: WorkflowSession
137125
): R = proceed(renderProps, renderState, null)
138126

139-
/**
140-
* Called after a workflow is rendered, and after the interceptor's [onRender] method. Every
141-
* call to this method is preceded by a call to [onBeforeRender].
142-
*
143-
* This method should be preferred over [onRender] since it supports additional types of workflows
144-
* that will not call [onRender].
145-
*/
146-
public fun <R> onAfterRender(rendering: R) {}
127+
@WorkflowExperimentalApi
128+
@WorkflowComposable
129+
@Composable
130+
public fun <P, O, R> onRenderComposeWorkflow(
131+
renderProps: P,
132+
emitOutput: (O) -> Unit,
133+
proceed: @WorkflowComposable @Composable (P, (O) -> Unit) -> R,
134+
session: WorkflowSession
135+
): R = proceed(renderProps, emitOutput)
147136

148137
/**
149138
* Intercept calls to [StatefulWorkflow.snapshotState] including the children calls.
@@ -399,35 +388,30 @@ internal fun <P, S, O, R> WorkflowInterceptor.intercept(
399388
renderProps: P,
400389
renderState: S,
401390
context: RenderContext<P, S, O>
402-
): R {
403-
onBeforeRender(renderProps, renderState, workflowSession)
404-
return onRender(
405-
renderProps,
406-
renderState,
407-
context,
408-
proceed = { props, state, interceptor ->
409-
// The `RenderContext` used *might* change - primarily in the case of our tests. E.g., The
410-
// `RenderTester` uses a special NoOp context to render twice to test for idempotency.
411-
// In order to support a changed render context but keep caching, we check to see if the
412-
// instance passed in has changed.
413-
if (cachedInterceptedRenderContext == null || canonicalRenderContext !== context ||
414-
canonicalRenderContextInterceptor != interceptor
415-
) {
416-
val interceptedRenderContext =
417-
interceptor?.let { InterceptedRenderContext(context, it) }
418-
?: context
419-
cachedInterceptedRenderContext = RenderContext(interceptedRenderContext, this)
420-
}
421-
canonicalRenderContext = context
422-
canonicalRenderContextInterceptor = interceptor
423-
// Use the intercepted RenderContext for rendering.
424-
workflow.render(props, state, cachedInterceptedRenderContext!!)
425-
},
426-
session = workflowSession,
427-
).also {
428-
onAfterRender(it)
429-
}
430-
}
391+
): R = onRender(
392+
renderProps,
393+
renderState,
394+
context,
395+
proceed = { props, state, interceptor ->
396+
// The `RenderContext` used *might* change - primarily in the case of our tests. E.g., The
397+
// `RenderTester` uses a special NoOp context to render twice to test for idempotency.
398+
// In order to support a changed render context but keep caching, we check to see if the
399+
// instance passed in has changed.
400+
if (cachedInterceptedRenderContext == null || canonicalRenderContext !== context ||
401+
canonicalRenderContextInterceptor != interceptor
402+
) {
403+
val interceptedRenderContext =
404+
interceptor?.let { InterceptedRenderContext(context, it) }
405+
?: context
406+
cachedInterceptedRenderContext = RenderContext(interceptedRenderContext, this)
407+
}
408+
canonicalRenderContext = context
409+
canonicalRenderContextInterceptor = interceptor
410+
// Use the intercepted RenderContext for rendering.
411+
workflow.render(props, state, cachedInterceptedRenderContext!!)
412+
},
413+
session = workflowSession,
414+
)
431415

432416
override fun snapshotState(state: S) =
433417
onSnapshotState(state, workflow::snapshotState, workflowSession)

‎workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/ChainedWorkflowInterceptor.kt

Lines changed: 49 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
package com.squareup.workflow1.internal
22

3+
import androidx.compose.runtime.Composable
34
import com.squareup.workflow1.BaseRenderContext
45
import com.squareup.workflow1.NoopWorkflowInterceptor
56
import com.squareup.workflow1.RenderingAndSnapshot
67
import com.squareup.workflow1.Snapshot
78
import com.squareup.workflow1.TreeSnapshot
89
import com.squareup.workflow1.Workflow
910
import com.squareup.workflow1.WorkflowAction
11+
import com.squareup.workflow1.WorkflowExperimentalApi
1012
import com.squareup.workflow1.WorkflowInterceptor
1113
import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
1214
import com.squareup.workflow1.WorkflowInterceptor.RuntimeLoopOutcome
@@ -75,16 +77,6 @@ internal class ChainedWorkflowInterceptor(
7577
return chainedProceed(renderProps)
7678
}
7779

78-
override fun <P, S> onBeforeRender(
79-
renderProps: P,
80-
renderState: S,
81-
session: WorkflowSession
82-
) {
83-
for (i in interceptors.indices) {
84-
interceptors[i].onBeforeRender(renderProps, renderState, session)
85-
}
86-
}
87-
8880
override fun <P, S, O, R> onRender(
8981
renderProps: P,
9082
renderState: S,
@@ -109,11 +101,54 @@ internal class ChainedWorkflowInterceptor(
109101
return chainedProceed(renderProps, renderState, null)
110102
}
111103

112-
override fun <R> onAfterRender(rendering: R) {
113-
// TODO does this still give optimized bytecode?
114-
for (i in interceptors.indices.reversed()) {
115-
interceptors[i].onAfterRender(rendering)
104+
@OptIn(WorkflowExperimentalApi::class)
105+
@Composable
106+
override fun <P, O, R> onRenderComposeWorkflow(
107+
renderProps: P,
108+
emitOutput: (O) -> Unit,
109+
proceed: @Composable (P, (O) -> Unit) -> R,
110+
session: WorkflowSession
111+
): R = onRenderComposeWorkflowStep(
112+
index = 0,
113+
stepProps = renderProps,
114+
stepEmitOutput = emitOutput,
115+
stepProceed = proceed,
116+
session = session
117+
)
118+
119+
/**
120+
* Recursive function for nesting chained interceptors' [onRenderComposeWorkflow] calls. We use
121+
* recursion for the compose call since it avoids creating a new list in the composition on every
122+
* render call.
123+
*/
124+
@OptIn(WorkflowExperimentalApi::class)
125+
@Composable
126+
private fun <P, O, R> onRenderComposeWorkflowStep(
127+
index: Int,
128+
stepProps: P,
129+
stepEmitOutput: (O) -> Unit,
130+
stepProceed: @Composable (P, (O) -> Unit) -> R,
131+
session: WorkflowSession
132+
): R {
133+
if (index >= interceptors.size) {
134+
return stepProceed(stepProps, stepEmitOutput)
116135
}
136+
137+
val interceptor = interceptors[index]
138+
return interceptor.onRenderComposeWorkflow(
139+
renderProps = stepProps,
140+
emitOutput = stepEmitOutput,
141+
proceed = { innerProps, innerEmitOutput ->
142+
onRenderComposeWorkflowStep(
143+
index = index + 1,
144+
stepProps = innerProps,
145+
stepEmitOutput = innerEmitOutput,
146+
stepProceed = stepProceed,
147+
session = session
148+
)
149+
},
150+
session = session
151+
)
117152
}
118153

119154
override fun onSnapshotStateWithChildren(

‎workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/ComposeWorkflowNode.kt

Lines changed: 13 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -203,13 +203,22 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
203203
// setContent will synchronously perform the first recomposition before returning, which is why
204204
// we leave cachedComposeWorkflow null for now: we don't want its produceRendering to be called
205205
// until we're actually doing a render pass.
206+
// We also need to set the composition content before calling startComposition so it doesn't
207+
// need to suspend to wait for it.
206208
composition.setContent {
207209
@Suppress("NAME_SHADOWING")
208210
val workflow = cachedComposeWorkflow
209211
if (workflow != null) {
210-
val rendering = workflow.produceRendering(
211-
props = lastProps,
212-
emitOutput = sendOutputToChannel
212+
val rendering = interceptor.onRenderComposeWorkflow(
213+
renderProps = lastProps,
214+
emitOutput = sendOutputToChannel,
215+
proceed = { innerProps, innerEmitOutput ->
216+
workflow.produceRendering(
217+
props = innerProps,
218+
emitOutput = innerEmitOutput
219+
)
220+
},
221+
session = this
213222
)
214223

215224
// lastRendering isn't snapshot state, so wait until the composition is applied to update
@@ -226,52 +235,10 @@ internal class ComposeWorkflowNode<PropsT, OutputT, RenderingT>(
226235
workflow: Workflow<PropsT, OutputT, RenderingT>,
227236
input: PropsT
228237
): RenderingT {
229-
230-
// Set the composition content before calling startComposition so it doesn't need to suspend to
231-
// wait for it.
232238
log("render setting content")
233239
this.cachedComposeWorkflow = workflow as ComposeWorkflow
240+
this.lastProps = input
234241

235-
return if (interceptor !== NoopWorkflowInterceptor) {
236-
// If there's an interceptor, then we have to stand up a bunch of fake stuff to pretend to be
237-
// a StatefulWorkflowNode.
238-
val renderContext = ComposeRenderContext<PropsT, OutputT>(
239-
runtimeConfig = runtimeConfig,
240-
workflowTracer = workflowTracer,
241-
actionSink = {
242-
// TODO(maybe): We _can_ actually support this, it's just gross. Need to make
243-
// outputsChannel support actions.
244-
throw UnsupportedOperationException("Actions not supported on ComposeWorkflows")
245-
}
246-
)
247-
248-
interceptor.onBeforeRender(
249-
renderProps = input,
250-
renderState = ComposeWorkflowState,
251-
session = this
252-
)
253-
interceptor.onRender(
254-
session = this,
255-
renderProps = input,
256-
renderState = ComposeWorkflowState,
257-
context = renderContext,
258-
proceed = { interceptedProps, _, _ ->
259-
this.lastProps = interceptedProps
260-
// TODO(maybe): Wire up calling back into the interceptor for outputs somehow? I think
261-
// this is more important than supporting sinkForInterceptor since it's used for logging.
262-
// But we need to be able to support sending actions to support this anyway.
263-
doRender()
264-
}
265-
).also {
266-
interceptor.onAfterRender(it)
267-
}
268-
} else {
269-
this.lastProps = input
270-
doRender()
271-
}
272-
}
273-
274-
private fun doRender(): RenderingT {
275242
val frameRequest = if (!lastRendering.isInitialized) {
276243
// Initial render kicks off the render loop. This should always synchronously request a frame.
277244
startComposition()

0 commit comments

Comments
 (0)