Enforcing Memory Allocation Patterns with an Effect System

Nov 23, 2014 by Todd Schiller

Writing performant Java software often means controlling allocation patterns to prevent garbage collection. For example, a program for processing financial FIX messages in real-time might use a fixed buffer to avoid allocations altogether.

In many cases, it's easy to determine whether or not code allocates memory. However, manually verifying allocation patterns becomes more difficult when method calls and method overriding are involved.

This post describes how to build an effect system to automatically enforce allocation patterns. If you've programmed in Java before, you've used an effect system: checked exceptions are an effect system.

Effect Systems

Unlike type systems, which indicate which type(s) an expression may have, effect systems indicate which effect(s) a statement or block may have. For example, the following method has the effect that it may throw a FileNotFoundException:

:::Java
public void maybeThrow() throws FileNotFoundException { ... }

The method above may throw the exception — or it may not. It's perfectly legal to annotate an empty method as potentially throwing an exception!

The compiler enforces that the caller of a method either: (1) handles the exception with a catch block, or (2) declares that it might also throw the exception (or a supertype of the exception).

Enforcing Allocation Patterns

For enforcing allocation patterns, we'd like to do something similar. The developer should be able to annotate which methods may allocate memory. Fortunately, the ability to define method annotations has existed since Java 1.5:

:::Java
@MayAlloc
public void mayAlloc(...) { ... }

More interestingly, the developer should be able to annotate which methods should not allocate memory:

:::Java
@NoAlloc
public void safeMethod(...) { ... }

Based on these annotations, we'd like the compiler to check that no method marked as @NoAlloc ever (1) itself allocates memory, or (2) calls a method that directly or indirectly allocates memory.

To make the compiler perform these checks, we'll use the Checker Framework, a tool for defining custom Java type systems. It's designed to take advantage of the new Type Annotations feature in Java 8. However, to implement a basic allocation effect system, we'll just use method annotations.

The full checker implementation is available on GitHub.

Defining Annotations

First we'll define the annotations @MayAlloc and @NoAlloc that the developer can write on method declarations:

:::Java
@Target(ElementType.METHOD)
public @interface MayAlloc { }

@Target(ElementType.METHOD)
public @interface NoAlloc { }

Since we're using the Checker Framework, we also have to define a type annotation for the default checker to process. For our checker, we'll just define a type annotation @EffectType that we'll tell the Checker Framework to apply everywhere:

:::Java
@TypeQualifier
@DefaultQualifierInHierarchy // Apply the annotation automatically
@SubtypeOf({}) // The annotation is not related to anything else
@ImplicitFor(trees = { Tree.Kind.NULL_LITERAL })
@Target({ ElementType.TYPE_USE, ElementType.TYPE_PARAMETER })
public @interface EffectType { }

For some effect systems, it is beneficial to define additional type annotations. See the Checker Framework's GUI Effect Checker as an example.

Creating a Checker

Now that we've defined our annotations, we can define a checker AllocEffect:

:::Java
@TypeQualifiers({ EffectType.class }) // Don't list the method annotations
public class AllocEffectChecker extends BaseTypeChecker { }

Every checker has two main parts. The Type Factory populates annotations on methods and types. The Visitor checks that the annotations are consistent, and that the code is consistent with the annotations. Using the Checker Framework naming convention, the type factory and visitor for our checker will be called AllocEffectTypeFactory and AllocEffectVisitor, respectively.

The Annotation Factory

The AllocEffectTypeFactory will compute annotations for methods that don't have developer-written annotations. Our factory will follow two rules:

If a method is unannotated, add the @MayAlloc annotation.
If an unannotated method overrides another method, inherit the annotation(s) from the overridden methods. If the developer writes a @MayAlloc annotation on a method overridding a method declared as @NoAlloc, emit a warning.

To enforce the second rule, we'll have the checker compute the range of effects each method can have based on the Java type hierarchy. Computing a range is overkill for an effect system with just 2 effects, but it provides a good basis for adding new features to the checker.

To define a range of effects, we have to define the relationship between each effect. Typically we'll use a hierarchy, just like in the standard Java type system. The two effects @NoAlloc and @MayAlloc have the following subtyping relation:

There's two ways to think about this. First, the set of methods that may or may not allocate memory must be a superset of the methods which may not allocate memory. Conversely, the @NoAlloc annotation provides a strictly stronger guarantee than the @MayAlloc annotation.

To check whether an operation (e.g., method call) is valid, we'll use a method checkEffect to check whether the operation's effect is a subtype of the effect of the method it's being called from. If not, we'll have Check Framework emit a warning:

:::Java
public void checkEffect(Effect callerEffect, Effect targetEffect, Tree node) {
  if (targetEffect.compareTo(callerEffect) > 0) {
    // The target effect is a supertype of the effect of the enclosing method

    // The message call.invalid.alloc is defined in a .properties file
    checker.report(Result.failure("call.invalid.alloc", targetEffect, callerEffect),
                   node);
  }
}

The Annotation Checker

Now that each method has an annotation, we'll define a checker AllocEffectVisitor to enforce the annotations. There are three rules we want to enforce:

If a method with a @NoAlloc annotation uses the new operator, emit a warning.
If a method with a @NoAlloc annotation calls a method with the @MayAlloc annotation, emit a warning.
If the developer has written both @MayAlloc and @NoAlloc on a method, emit a warning.

Using these rules, it's always safe for a method to call another method that has the @NoAlloc annotation.

The Checker Framework's visitor implementation, BaseTypeVisitor, recursively visits each node in the program's Abstract Syntax Tree. To enforce the first rule, we'll override the visitNewArray and checkConstructorInvocation visitor methods. Our implementation will check that the enclosing method has the @MayAlloc annotation:

:::Java
@Override
public Void visitNewArray(NewArrayTree node, Void p) {

  // The 'new' operator is just like a method that has a @MayAlloc annotation
  Effect targetEffect = new Effect(MayAlloc.class);

  // Get the annotation for the enclosing method
  MethodTree callerTree = TreeUtils.enclosingMethod(getCurrentPath());
  ExecutableElement callerElt = TreeUtils.elementFromDeclaration(callerTree);

  // Make sure the two effects are consistent. That is, that target
  // effect is a subtype of the calling effect. Because the
  // targetEffect is @MayAlloc, this checks that the method
  // annotation is @MayAlloc too.
  checkEffect(callerEffect, targetEffect, node);

  return super.visitNewArray(node, p);
}

To enforce the second rule, we'll override the visitMethodInvocation method. The implementation is the same, except we use the effect of the method being called:

:::Java
@Override
public Void visitMethodInvocation(MethodInvocationTree node, Void p) {

  // Get the annotation for the method being called
  ExecutableElement methodElt = TreeUtils.elementFromUse(node);
  Effect targetEffect = atypeFactory.getDeclaredEffect(methodElt);

  // ... same as for visitNewArray(...)
}

To enforce the third rule, we'll override the visitMethod method to report an error if a method has both a @MayAlloc and @NoAlloc annotation:

:::Java
@Override
public Void visitMethod(MethodTree node, Void p) {
  ExecutableElement methElt = TreeUtils.elementFromDeclaration(node);
  AnnotationMirror mayAlloc = atypeFactory.getDeclAnnotation(methElt, MayAlloc.class);
  AnnotationMirror noAlloc = atypeFactory.getDeclAnnotation(methElt, NoAlloc.class);

  if (mayAlloc != null && noAlloc != null) {
    checker.report(Result.failure("annotations.conflicts"), node);
  }

  return super.visitMethod(node, p);
}

Running the Checker

At this point, we've built a basic checker for memory allocations. You can download the full implementation on GitHub. To run the checker, use javac with the -processor argument and the checker on the classpath (setting the classpath with the -cp flag, as necessary):

:::bash
javac -processor com.toddschiller.checker.AllocEffectChecker com/toddschiller/experiments/AllocationEffects.java

The implementation on GitHub also supplies a debug mode, which you can enable by passing the -Alint=debugSpew flag to the compiler. The debug mode reports each check performed by the checker.

For the example file AllocationEffects.java, the checker reports 6 warnings. The two assignments

:::Java
@NoAlloc
public void shouldWarn() {
  int x = mayAllocateMemory(); // an unannotated method
  int y = new Integer(3);
}

are both flagged by the checker:

AllocationEffects.java:42: error: [call.invalid.alloc] Calling a method with MayAlloc effect from a context limited to NoAlloc effects.
            int i = new Integer(2);
                    ^
AllocationEffects.java:56: error: [call.invalid.alloc] Calling a method with MayAlloc effect from a context limited to NoAlloc effects.
        int x = mayAllocateMemory();
                                 ^

Improving the Checker

Now that we've implemented a basic checker, there's two important questions to ask:

Can the checker emit a warning even when the program won't allocate memory (false positives)?
More importantly, are there any allocations that the checker will fail to report (false negatives)?

Eliminating False Warnings

The first question is easier to answer — of course! Just like the Java type checker, our checker doesn't consider values or control flow:

:::Java
public void conditionallyAlloc(boolean alloc) {
  if (alloc) doAlloc();
}

@NoAlloc
public void noAlloc(){
  conditionallyAlloc(false); // WARNING!
}

Looking at the implementation of conditionallyAlloc, we ideally wouldn't get a warning. But, if we were to write an annotation for the the conditionallyAlloc method, what would we write? Unfortunately our vocabulary of just @NoAlloc and @MayAlloc is insufficient.

One way to solve the problem would be to let the developer write an annotation @MayAllocIf("alloc"), referring to the boolean alloc parameter. While we could modify the checker to support this particular annotation, we quickly run into a wall trying to support other expressions — we'd have to figure out how to verify arbitrary logic at compile-time!

In this case, a reasonable solution might be to ignore the warning by adding an @SuppressWarnings("alloceffect") annotation to the call to conditionallyAlloc. ("alloceffect" is the name of the checker in lower-case).

However, in many cases, it is better to refactor the code. Code that a checker can reason about is generally also easier for a human developer to reason about.

Can you think of other cases for which the checker will issue false warnings? Are there common patterns (e.g., the singleton pattern) where it would make sense to introduce special annotations?

Identifying All Allocations

Identifying allocations that the checker will miss is trickier — are all allocations the result of the new keyword?

Ignoring reflection and system calls, the answer appears to be "yes". However, some uses of new can be obfuscated. Consider a class with a static field and a method that accesses that field:

:::Java
public class AllocOnLoad {
  public final static Integer five;
  static {
    // The static initializer can contain arbitrary code
    five = new Integer(5);
  }
}

public class Other {
  @NoAlloc
  public void addFive(int x){
    return x + AllocOnLoad.five; // Warning?
  }
}

When the program runs, if a call to addFive is the first time the AllocOnLoad class if referenced during the execution, the static initializer will run and memory will be allocated. On all subsequent calls, no memory will be allocated.

To catch this behavior, the developer would need annotate classes with whether or not their static initializer may allocate memory. The checker could then treat class references similar to method calls.

The problem with this approach is that it would lead to a lot of false positives. A potential solution would be to allow the developer to specify that certain classes have already been loaded when a method is called, à la:

:::Java
@NoAlloc
@RequireLoaded({AllocOnLoad.class})
public void addFive(int x) {
  return x + AllocOnLoad.five; // Safe!
}

Implementing such a system would involve introducing an effect system to track which classes each method loads. The difference is that the effect system would need to track which classes must have been loaded. For example, the addFive method would have the effect @Loads({AllocOnLoad.class}) because the method always references the AllocOnLoad class.

Writing down which classes a method must load is tedious and error-prone. The developer has to consider all the classes transitively referenced by the method. Therefore, this information would be a good candidate for automatic inference.

Can you think of other cases where the checker would miss an allocation? What happens when you concatenate two strings: str1 + str2?

Summary

This post described an effect system for enforcing memory allocation patterns. We introduced two method annotations, @MayAlloc and @NoAlloc, to describe the effect of each method. We then used the Checker Framework to enforce that no memory allocations are either directly or indirectly performed by methods marked as @NoAlloc.

For more information about implementing effect systems with the Checker Framework, check out the framework's GUI Effect Checker, which checks for UI threading errors. Unlike the basic checker we built, it makes use of the Checker Framework's type annotation features to support common UI programming idioms.

The source code for the allocation checker is available on GitHub at: https://github.com/twschiller/alloc-effect-checker

Todd Schiller

Machine ✘ Human Intelligence