Automatic Feature Flag Cleanup with Codemod

Introduction

When a feature flag becomes stale, developers need to manually remove it from the codebase. This process involves more than just replacing the feature flag SDK call with a static value; it also requires a thorough refactoring of the code to ensure it remains clean after the change. The task can be intricate, especially in larger codebases where feature flags might be deeply integrated into various components and logic flows.

To illustrate this, let's look at an example of such a refactoring process. The code snippet from the before state, which includes conditional logic based on the feature flag, can be significantly simplified in the after state once the flag is removed and the code is refactored.

Before

const myFlag = useVariableValue('feature-name', false);

if(myFlag) {
	console.log('flag is true');
}

if(!myFlag) {
	console.log('flag is false');
}

After

console.log('flag is false');

To automate this process, we built codemods that autonomously handles both the removal of feature flags and the subsequent refactoring of the code. These codemods significantly streamline the cleanup process by automating the identification and replacement of feature flag SDK calls with static values, as well as performing the necessary code refactoring to maintain code quality.

Currently we added codemods for multiple popular feature flag providers, such as DevCycle and Statsig, but in future we can easily add support for many other providers.

Architecture overview

To make it easy to support multiple feature flag providers, we split the code into two main components:

1. Feature Flag SDK Specific Code (Provider): This component encapsulates SDK-specific knowledge, such as SDK method names, import names, and method signatures. Essentially, it understands how to match SDK calls and determine the appropriate values to replace them with. For each SDK, we implement two methods: getMatcher and getReplacer. The getMatcher method is responsible for identifying the correct CallExpression, while the getReplacer method returns the appropriate replacement node.
2. Unused Code Removal: This component handles the removal of code that is no longer needed due to the feature flag cleanup. It is designed to be generic, allowing easily add new providers. Unused code removal includes multiple atomic steps that are executed one by one. This diagram shows the top level architecture of the codemod.

Modular & customizable design for feature flag cleanup codemod

Unused code removal

Let's take a closer look at each step of code removal component.

Simplify Object Properties

This step involves removing redundant member expressions that remain after replacing the feature flag SDK call with an object literal.

Before

({ value: 'theValue' }).value

After

'theValue'

Update References

If flag is used as the part of variable declaration, we want to replace the variable associated with the feature flag as well.

Before

const var1 = useFlag(user, 'simple-case', true);
const var2 = var1;

console.log(var1);

if(var1) {
 console.log('var2 is true');
}

After

const var3 = true;

console.log(true);

if(true) {
 console.log('var2 is true');
}

Simplify Unary Expressions

This stepis aiming for simplifying unary expression with exclamation mark:

This transformation is only applied to the literals.

Before

!!!false 
!!"string"
""

After

true
true
false

Simplify Binary Expressions

This step replaces useless binary expressions. When left and right values are literals we can evaluate such expressions and replace the expression with its result.

Before

true === true
1 === true
```

After

true
false

Simplify Logical Expressions

This step is similar with previous, here we try to simplify useless logical expressions. In javascript logical expression with ampersand returns the first falsy value, if we see expression that starts with true or truthy literal, we can just remove the literal. When expression starts with the falsy literal, replace the expression with this literal.

Before

true && x 
false && x

After

x
false

Remove Unnecessary Parenthesis

Sometimes after simplifying nested expression we have useless parenthesis, this step removes them.

Simplify If Statements

This step simply unwraps the if statement if condition is truthy, or removes it if its falsy. On this step we eliminate useless branching.

Before

if(true) {
 console.log('truthy')
}

if(false) {
 console.log('truthy')
}

After

console.log('truthy')

Simplify Conditional Expressions

This step is similar as previous, but for ternary expressions:

Before

true ? 'truthy' : 'falsy'

After

'truthy'

To handle nesting, we run all steps multiple times.

Adding support for custom feature flag provider

Adding support for a new provider is quite straightforward. We simply need to define the getMatcher and getReplacer methods.

getMatcher accepts the feature flag key name and return the matcher function that accepts the CallExpression and should return the SDK method name or undefined if we want to skip the node.

getReplacer accepts the key, type and the value that should be used for replacement of the feature flag. Additionally it accepts the SDK method name, because we can have different replacements for different SDK method calls.

getReplacer should return the node that will be stringified and used as replacement for feature flag call in the main part of the codemod.

Here is simple example of the custom provider that matches method with the name useFlag and given feature flag key name and replaces it with the literal.

const SimpleProvider = {
  getMatcher: (keyName: string) => (ce: CallExpression) => {
    const name = getCEExpressionName(ce);
    // only match useFlag call expression
    if (name !== "useFlag") {
      return null;
    }
    
    const args = ce.getArguments();
    
    const keyArg = args[0];
    
    // only match the useFlag with given keyName
    if(!Node.isStringLiteral(keyArg) || keyArg.getLiteralText() !== keyName) {
     return;
    }

    return { name };
  },
  // replacer accepts key, type (because we need to cast the value that is passed as string)
  // value and the matched callExpressions name 
  getReplacer: (
    key: string,
    type: VariableType,
    value: VariableValue,
    { name }: string,
  ) => {
    return buildLiteral(type, value);
  },
};

Conclusion

Cleaning up feature flags is a major hassle; developers in large teams waste months on manual work. Fortunately, this can be solved with a customized codemod, a practice established in big tech for years. We've developed a modular and customizable codemod so that anyone can easily adapt it to their needs. To further customize our codemods, you can use Codemod Studio for AI assistance or contact us to build professional-grade codemods tailored to your specific requirements. Codemod automates repetitive code maintenance tasks for feature flags, reducing friction for experimentation and leading to more and better features for end users.