Refactoring Code with AOP

Now that we know the theory and can apply AOP in practice, let's make our code better.

"Break everything and rewrite from scratch" — that's not our style! Refactoring is about careful improvements that don't change behavior for the user. AOP fits nicely here.

How do you know it's time to use AOP? Look at your code. Do you see identical logging snippets sprinkled around? Lots of access checks? Similar error handling? Those are clear signs it's time to pull repeated code into aspects.

Analyzing problems in the code

Say you have a service that manages users. Here's an example:

@Service
public class UserService {

    private final UserRepository userRepository;

    public UserService(UserRepository userRepository) {
        this.userRepository = userRepository;
    }

    public User createUser(User user) {
        System.out.println("Starting user creation: " + user.getName());
        try {
            User savedUser = userRepository.save(user);
            System.out.println("User created successfully: " + savedUser.getName());
            return savedUser;
        } catch (Exception e) {
            System.err.println("Error while creating user: " + e.getMessage());
            throw e;
        }
    }

    public User getUser(int id) {
        System.out.println("Requesting user with ID: " + id);
        return userRepository.findById(id)
                .orElseThrow(() -> new RuntimeException("User with that ID not found"));
    }
}

As you can see, business logic (creating and fetching users) is mixed with cross-cutting concerns (logging). Also, exception handling is duplicated across methods. It looks messy and is hard to maintain.

Extracting logging into an aspect

The first step in refactoring is to extract logging into an aspect. We'll create an aspect that logs method start and finish, and also handles exceptions.

Step 1: Create a logging aspect

@Aspect
@Component
public class LoggingAspect {

    private static final Logger logger = LoggerFactory.getLogger(LoggingAspect.class);

    @Around("execution(* com.example.service.*.*(..))")  // Apply to all methods in the service package
    public Object logAroundMethods(ProceedingJoinPoint joinPoint) throws Throwable {
        String methodName = joinPoint.getSignature().getName();
        logger.info("Method start: {}", methodName);

        Object result;

        try {
            result = joinPoint.proceed(); // Execute the target method
            logger.info("Method completed successfully: {}", methodName);
        } catch (Exception e) {
            logger.error("Error while executing method: {}", methodName, e);
            throw e; // Make sure to rethrow the exception!
        }

        return result;
    }
}

With the @Around annotation we create aspects that wrap the execution of the target method. Now the logging logic is completely separated from business code, and the service code becomes a lot cleaner.

Updating the service after refactoring

Now our UserService looks like this:

@Service
public class UserService {

    private final UserRepository userRepository;

    public UserService(UserRepository userRepository) {
        this.userRepository = userRepository;
    }

    public User createUser(User user) {
        return userRepository.save(user);
    }

    public User getUser(int id) {
        return userRepository.findById(id)
                .orElseThrow(() -> new RuntimeException("User with that ID not found"));
    }
}

As you can see, the logic is now more compact and cleaner. Logging now lives in the aspect, and we can use it across all services without changing their code.

Optimizing exception handling

The next step is exception handling. Right now exceptions are partly handled in the logging aspect, but often developers want centralized control over error handling.

Step 2: Create an error-handling aspect

@Aspect
@Component
public class ErrorHandlingAspect {

    private static final Logger logger = LoggerFactory.getLogger(ErrorHandlingAspect.class);

    @AfterThrowing(pointcut = "execution(* com.example.service.*.*(..))", throwing = "ex")
    public void handleException(Exception ex) {
        // Log the error
        logger.error("An exception occurred: {}", ex.getMessage(), ex);
        // You can add extra handling here, e.g., notify the team
    }
}

This aspect will trigger after exceptions are thrown in services. It takes care of logging errors without cluttering the main code.

Advanced use: security checks

Say we want to restrict access to methods based on the user's role. For example, only admins can create users.

Step 3: Create a security-check aspect

@Aspect
@Component
public class SecurityAspect {

    @Before("execution(* com.example.service.UserService.createUser(..))")
    public void checkCreateUserAccess() {
        // Check if the current user has the right to create users
        // The code below is just an example!
        boolean hasAccess = SecurityContextHolder.getContext().getAuthentication().getAuthorities()
                .contains(new SimpleGrantedAuthority("ROLE_ADMIN"));

        if (!hasAccess) {
            throw new SecurityException("The current user doesn't have permission to perform this action");
        }
    }
}

Now access to createUser will be checked automatically before it's invoked.

Testing the changes

After refactoring it's important to make sure we didn't break anything. Add tests to verify:

1. Logging works and contains correct entries.
2. Exceptions are handled properly and the app doesn't crash.
3. Security checks correctly restrict access.

Helpful tips and common mistakes

• Too many aspects: don't create aspects for every tiny thing, otherwise AOP becomes overkill and complicates the project.
• Performance: AOP can add some overhead if used on very frequently called methods. For example, avoid logging in aspects for methods executed inside high-frequency loops.
• Mistakes in pointcut expressions: misconfigured pointcut expressions can accidentally intercept extra methods or, conversely, ignore important ones. Test your expressions carefully.
• Proxy gotchas: remember that aspects work via proxies. If you call a method from the same class where it's defined, the aspect might not trigger.

Wins after refactoring

After refactoring with AOP the code became:

1. Clean: service logic no longer contains duplicated code for logging, security, or exception handling.
2. Modular: cross-cutting concerns are moved into aspects, which makes them easier to maintain.
3. Easy to change: we can add or modify behavior, like logging or security, by changing only the aspects instead of every service method.

That's AOP magic! It's like a thin layer between your business logic and "everything else". Next time someone says their code is "perfect without aspects", show them our refactor. 😉