Subtle retain cycle is subtle

Retain cycles are often a pain to track down and here’s an example of a less obvious one we had recently.

The problem

Here’s a simplified reproduction of the issue

class Example {
    var task: Task<Void, Never>?

    init() {
        task = Task { [weak self] in
            guard let self else { return }

            repeat {
                performSomeWork()
            } while !Task.isCancelled
        }
    }

    func performSomeWork() { }

    deinit {
        print("deinit")
        task?.cancel()
    }
}

Let’s not focus too much on the exact code as it doesn’t do anything except illustrate the issue. When running this code the deinit will never run because the Task is creating a retain cycle keeping Example alive indefinitely.

At first glance this looks like fairly standard code - the part of interest is

task = Task { [weak self] in
    guard let self else { return }

    repeat {
        performSomeWork()
    } while !Task.isCancelled
}

In the above we see the common weak/strong dance that we are all used to but we still have a cycle so what gives?

We are spinning a loop in the task that only stops when the task is cancelled. The only place we currently call cancel is in the deinit of the Example class so this loop is partly responsible for the cycle. The key thing to look for is who is taking a strong reference and what is the scope of that reference?

task = Task { [weak self] in        //
    guard let self else { return }  // - strong reference taken here
                                    //
    repeat {                        //
        performSomeWork()           //
    } while !Task.isCancelled       //
}                                   // - goes out of scope here

The problem we have looking at the scope is that the strong reference is in scope until the end of the function, but we have our repeat loop before the end of the function so we will never get to the end.

Breaking the cycle

There’s many ways to break the cycle - let’s look at a few

Change the scope of the strong reference

task = Task { [weak self] in            //
    repeat {                            //
        guard let self else { return }  // - strong reference taken here
        performSomeWork()               //
    } while !Task.isCancelled           // - goes out of scope here
}                                       //

If we move the guard inside the repeat then it will only take a strong reference for the body of repeat. This means that the strong reference is released and retaken each time through the loop. Due to the guard being evaluated fresh each time this allows the cycle to be broken.

Use the weak reference everywhere

task = Task { [weak self] in
    repeat {
        self?.performSomeWork()
    } while !Task.isCancelled
}

In this example it looks pretty clean to do this but in real code you might not be able to have the nullability in which case you’d end up using guard or if let to unwrap things (just be careful on scope).

Manually break the cycle

task?.cancel()

For this you’d have to have some other code get a reference to the task and call cancel() at the appropriate time.

Be careful

Another thing you might try to break the cycle is using the capture groups.

task = Task { [performSomeWork] in
    repeat {
        performSomeWork()
    } while !Task.isCancelled
}

For this example we are back to retain cycle city. The issue is instance methods have an implicit reference to self so this won’t do the job.

The capture group would indeed work if we are getting a reference to something that doesn’t have a self reference for example instance properties.

You could write a unit to verify that the reference does not leak something like this. In this example though you’d need to add a short delay before you set the system under test to nil to ensure that the Task has had enough time to start working and take any strong reference you want to validate is held correctly.

Conclusion

Retain cycles are a pain and the ol’ chuck a weak on things doesn’t always work so it’s worth writing tests and using instruments to hunt things down.

KSP and Me

I’ve been using Kotlin Symbol Processing (KSP) for a few years so I thought I’d reflect on how I like to work with it to stay productive.

First things First

Let’s start by recognising if you are new to KSP it is hard to get up to speed, it’s not impossible but it will require some graft to really get stuck in. Many of the blog posts I read when I was starting were very good at helping you get something compiling but then pretty much finished there. Without someone holding my hand or giving me cues of where to look I was kind of stuck not really knowing the potential of the tool I was learning.

Don’t treat what you read on the internet as gospel

Many of the blog posts I read when starting out had a similar pattern of suggesting you should use the visitor pattern and KotlinPoet, without really saying why you’d want to use them. I’ve read the Gang of Four book many moons ago but had all but forgotten the visitor pattern and I’d never heard of KotlinPoet so that’s two things I was expected to learn just to follow an introductory tutorial.

Thankfully I’m a few years in and I’ve mostly managed to avoid using the visitor pattern for my use cases. My coding style these days leans more towards a functional style so less common OO patterns just feel alien and slow me down.

For example to get all of a class’ nested children I could use the visitor pattern something like this:

class MyVisitor: KSDefaultVisitor<Unit, Klass?>() {
    override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) = Klass(
        classDeclaration.simpleName.asString(),
        classDeclaration.declarations.mapNotNull { it.accept(this, Unit) }.toList()
    )

    override fun defaultHandler(node: KSNode, data: Unit): Klass? = null
}

Tangent: None of the examples I read at the time actually accumulated results in this functional style using the second type parameter but instead opted but having an instance variable that you accessed after parsing was completed.

Or I could use a more functional style like this:

fun KSClassDeclaration.nestedClasses(): Klass = Klass(
    simpleName.asString(),
    declarations.filterIsInstance<KSClassDeclaration>().map(KSClassDeclaration::nestedClasses).toList()
)

The functional style I personally find more direct and I can see the recursion happening I’m not relying on learning what the various conformances to visitor are and which is right for my use case and the methods I need to use/implement (accept, defaultHandler) and why.

Anyway I’m not trying to sell one approach over the other because that’s for you and your team to thrash out. I’m mostly just saying if it works then use it, you don’t have to feel like I did that I was somehow holding it wrong because my code didn’t look like all the blog posts I was reading.

The other good thing to report is that I haven’t needed to learn KotlinPoet, again for the things I’ve worked on multiline string literals have been more than adequate. I mean I know what the Kotlin code I want to generate should look like so having an extra layer in the middle doesn’t add much for me personally.

Separate parsing and generating

When I started I kept trying to build up the final String of what the Kotlin source code should be whilst parsing the code. This is not a great idea as you soon tie yourself in knots. What compilers tend to do, which is the pattern I follow now is

• Parse the input into some intermediate representation
• Process the intermediate representation
• Render the output

For step 1 I like to take the types provided by KSP such as KSClassDeclaration and extract out the information I need into simple data class types. That way the processing logic I write next doesn’t need to know about KSP and the task is more focussed on gathering all the data that my processor thinks is interesting.

Once I have the data I’ll then do any validation, filtering or mapping to new representations. At this point I’m working with simple immutable data classs with well named properties, which is much preferred to having all my business logic calling all combinations of deeply nested resolve(), declaration, asString(), etc.

The final step is rendering, which is very often now just a collection of string templates that interpolate in the nicely structured data from the previous step.

I think there are a few great advantages to separating things out:

• You can generate different code for different targets (e.g. Kotlin/JS vs Kotlin/JVM) in a much simpler way
• Future readers don’t have to follow a potential mess of building a string whilst parsing
• More easily add unit tests around the business rules in the processor

Validate before you generate

Linked to the mistake mentioned in the section above about trying to do things in one go I would fully recommend writing out the code you want to generate manually and checking it works. I’ve found that I was constantly starting off with a simple picture of what I needed to generate and it seemed so obvious what was needed that I started writing the generation code. The issue is I’m not a very good developer and the simple code I imagine never really works and often requires changes. It’s much simpler to edit, compile and run code directly rather than trying to change the code to generate new code so that you can run and validate it.

Example use cases

The biggest pain point for me was not having that spark of inspiration for what I could be doing with KSP. Here’s a few things that me and my team have used KSP for:

• Validation tasks
  • Ensuring that a module correctly uses functions or computed properties, this is a bit niche but this module houses UI strings and if we use properties then every single string would end up in the JS artifact even if it wasn’t referenced. To avoid every JS artifact having every possible string we rely on the dead code elimination you get when the compiler notices you don’t invoke a function.
  • Ensuring that certain types conform to Serializable to support Android. If we forget the conformance then we could crash at runtime if an activity tries to serialize state.
• Generate type aliases
  • KMP doesn’t export typealises to iOS and the naming rules for types can be a little funky. Some times subtypes have dot separators (Parent.Child) and other times the symbol names are just smashed together (ParentChild). This is super confusing and we want to alias the most recent versions of some generated types so iOS developers never know about the actual versioning. The processor for this outputs Swift code, which is then packaged via SKIE.
• Generate type safe routing helpers
  • A colleague wrote a processor that will read various spring boot annotations to calculate the path for an endpoint and what arguments are required. This is then all used to generate typesafe helper functions that allow people to do routing in a much safer manor.
• Generate DSL versions
  • Me and a colleague wrote a pretty comprehensive processor that generates type safe versioned DSLs allowing us to migrate away from a system that meant adding a new version of our DSL required fairly specialised knowledge of our versioning system + many hours-days of work and resulted in inconsistent results to now mostly just bumping a number.
• Generating observability wrappers
  • Me and various colleagues wrote a processor that takes a class with some annotations sprinkled on it and generates a wrapper class that knows when properties are being written and will require us to recompute a new state. This processor also generates type safe bindings that allow us to bind our UI to these properties.
• Generate per request caches in webflux
  • I’m still learning webflux but a requirement came up to have a per request cache, this would normally be done with an @Cacheable annotation on a method in a normal thread based spring boot application. What I ended up spiking was having KSP look for an annotation of @Memoize which then generated a CoWebFilter to create a typed caffeine cache and slap it in the coroutine context. Then the KSP would generate a wrapper class that delegates to the original after trying the cache. This generated delegate wrapper would have a @Primary annotation so spring would wire it in rather than original.

There’s plenty more example uses out there these days if you look around but all of the above are either in live active projects or hopefully will be soon.

Conclusion

I think it’s great to have good documentation on how to use a library but sometimes the thing that is missing is the little bit of inspiration that get you thinking about how you could apply a technology to your project. I’m glad we embraced KSP and we have done away with so much boilerplate code and all the opportunities for mistakes and inconsistencies to sneak in that makes maintenance harder.

CustomTestStringConvertible

SwiftTesting’s parameterised tests are really great. I’ve been finding that I often want to give each example a nice name that shows in the test navigator as the default String(describing:) approach doesn’t always hit the mark.

Let’s take the contrived example of testing an enum initialiser that allows the enum to be constructed with a case insensitive string.

Attempt 1

As a first stab we might go for

struct SoftwareTests {
    @Test("Lowercase names can be used")
    func lowercaseNamesCanBeUsed() {
        #expect(Software("macos") == .macOS)
    }

    @Test("Uppercase names can be used")
    func uppercaseNamesCanBeUsed() {
        #expect(Software("MACOS") == .macOS)
    }

    @Test("Unknown names are mapped to unknown")
    func unknownNamesAreMappedToUnknown() {
        #expect(Software("??") == .unknown)
    }
}

This is fine and results in the test navigator showing our tests like this:

└── SoftwareTests
    ├── Lowercase names can be used
    ├── Uppercase names can be used
    └── Unknown names are mapped to unknown

This all looks fairly reasonable but even in this simple example we can see duplication. Each test repeats the exact same pattern in the expectation. Full disclosure I probably wouldn’t bother changing this code as it’s already fairly concise but let’s imagine that the test bodies were a little longer and there was duplicated set up happening in the bodies.

Attempt 2

In this case you’d want to jump to some parameterised tests which might look like this:

struct SoftwareTests {
    @Test(
        "Init is case insensitive and handles unknown cases",
        arguments: [
            (input: "macos", expected: Software.macOS),
            (input: "MACOS", expected: Software.macOS),
            (input: "??", expected: Software.unknown),
        ]
    )
    func initIsCaseInsensitiveAndHandlesUnknownCases(input: String, expected: Software) {
        #expect(Software(input) == expected)
    }
}

The duplication is gone and the different permutations are run in parallel which is great for test performance. The issue is that we’ve made the test navigator view a little less useful as it now looks like this:

└── SoftwareTests
    └── Init is case insensitive and handles unknown cases
        ├── "macos", .macOS
        ├── "MACOS", .macOS
        └── "??", .unknown

Those labels don’t really mean much unless you read the test implementation. Something to note is that the actual arguments declaration in the @Test annotation is using labels to make it easier to read the test set up to know which field is the input vs the expected. Although the code source is enhanced with these labels the test navigator is not so clear.

Attempt 3

Let’s fix the previous issue using the CustomTestStringConvertible protocol

struct SoftwareTests {
    @Test(
        "Init is case insensitive and handles unknown cases",
        arguments: CaseInsensitiveInit.allCases
    )
    func initIsCaseInsensitiveAndHandlesUnknownCases(fixture: CaseInsensitiveInit) {
        #expect(Software(fixture.input) == fixture.expected)
    }

    struct CaseInsensitiveInit: CustomTestStringConvertible, CaseIterable {
        let input: String
        let expected: Software
        let testDescription: String

        static let allCases: [CaseInsensitiveInit] = [
            .init(input: "macos", expected: .macOS, testDescription: "Lowercase names can be used"),
            .init(input: "MACOS", expected: .macOS, testDescription: "Uppercase names can be used"),
            .init(input: "??", expected: .unknown, testDescription: "Unknown names are mapped to unknown"),
        ]
    }
}

With this set up the test navigator is looking much nicer:

└── SoftwareTests
    └── Init is case insensitive and handles unknown cases
        ├── Lowercase names can be used
        ├── Uppercase names can be used
        └── Unknown names are mapped to unknown

We’ve restored the handy naming whilst keeping the ability for the framework to optimise and call all the cases in parallel.

Going further

With the above we have to add boiler plate but the benefits are quite useful. There are common types where we can provide helper functions to make this process a little smoother like booleans. If we create a little helper like this:

struct DescribedBool: CustomTestStringConvertible {
    let testDescription: String
    let value: Bool
}

func boolTestStringConvertible(label: (Bool) -> String) -> [DescribedBool] {
   [
    .init(testDescription: label(false), value: false),
    .init(testDescription: label(true), value: true),
   ]
}

Then we can write tests that can be described a lot easier

@Test("Display an appropriate install state", arguments: boolTestStringConvertible { "installed = \($0)"})
func displayAnAppropriateInstallState(fixture: DescribedBool) {
    // ...
}

NB: The above will hopefully work in future but due to way the macro currently works it doesn’t like there being a closure in the arguments position. We can work around this by adding a little shim

@Test("Display an appropriate install state", arguments: installedCases())
func displayAnAppropriateInstallState(fixture: DescribedBool) {
    // ...
}

private static func installedCases() -> [DescribedBool] {
    boolTestStringConvertible { "installed = \($0)" }
}

With this in place we get nice descriptions like this:

└── SoftwareTests
    └── Display an appropriate install state
        ├── installed = false
        └── installed = true

Conclusion

SwiftTesting parameterised tests are great. It’s also very easy just to slap a load of test cases in simple tuples and exercise a lot of things but maybe lose some clarity in the test navigator around what the tests are doing. Using CustomTestStringConvertible is a nice way to bring some order back and help yourself and other travellers of your code base to navigate some hopefully extensive tests suites.