DEV Community: Anuj Singh

Enhancing Transparency and Accountability: Implementing Entity Audit Logging in Java [Part-2/2]

Anuj Singh — Fri, 28 Jun 2024 11:01:41 +0000

In the last post, we used io.ebean's ChangeLog annotation to log entity changes in a synchronous blocking I/O application. In this post, we will see how that works in a multithreaded or async (non-blocking) I/O application.

If your system uses a separate thread-pool or executor context for doing I/O tasks such as making an HTTP API request or a database transaction, then, the audit log prepare class will not be able to detect the user context on model entities that are updated, as the model entity will be saved in a different context or thread-pool. Hence, we need to ensure that the thread-context is not lost when switching the executor or when switching threads in the same thread-pool.

To do so, first we need to understand how to store some data in thread context. For this we will use ThreadLocals. ThreadLocal construct allows us to store data that will be accessible only by a specific thread. Each thread will have its own ThreadLocal instance, hence while switching the context, we need to ensure that the thread context of the current thread is propagated to the next thread.

STORING USER CONTEXT IN THREAD LOCALS

We can set ThreadLocals with user info and use it in change log prepare’s implementation to set it in change set.

ThreadLocalManager

public class ThreadLocalManager {
    private static ThreadLocal<Map<String, Object>> context = new ThreadLocal<>();
    public static void addToContext(String key, Object value) {
        Map<String, Object> currentContext = ThreadLocalManager.context.get() == null ? new HashMap<>() : context.get();
        currentContext.put(key, value);
        context.set(currentContext);
    }
    public static void setContext(Map<String, Object> contextMap) {
        context.set(contextMap);
    }
    public static Map<String, Object> getContext() {
        return context.get();
    }
    public static Object getFromContext(String key) {
        return Nullifier.get(() -> context.get().getOrDefault(key, "NA"));
    }
    public static void clearContext() {
        ThreadLocalManager.context.remove();
    }
}

We can add additional info before setting the thread local

ThreadLocalManager.addToContext("userContext", new HashMap<String, String>() {{
    put("userName", "some user");
    put("userEmail", "some.user@company.com");
}}
entity.save()

Now we'll modify our ChangeLogPrepare to read user data from thread context

Custom ChangeLogPrepare

public class AuditLogPrepare implements ChangeLogPrepare {
    private final play.Logger.ALogger logger = Logger.of(this.getClass());
    @Override
    public boolean prepare(ChangeSet changes) {
        Map<String, String> userContext = Nullifier.get(() -> (Map<String, String>) ThreadLocalManager.getContext().get("userContext"), new HashMap<>());
        if (userContext.isEmpty()) logger.warn("[ALERT] userContext is empty for changeset: " + changes.toString());
        changes.getUserContext().put("userName", authMap.getOrDefault("userName", DEFAULT_USER_NAME));
        changes.getUserContext().put("userEmail", authMap.getOrDefault("userEmail", DEFAULT_USER_EMAIL));
        changes.setSource("MyApp");
        return true;
    }
}

As we see, the user data is taken from thread local, we need to ensure that the thread context is maintained while switching the thread. For that, we'll create a utility class that helps us propagate the thread context to next runnable/callable.

ContextUtility

public class ContextUtility {
    public static <T> Callable<T> wrapWithContext(Callable<T> task) {
        Map<String, Object> previousContext = ThreadLocalManager.getContext();
        if (previousContext == null)
            return task;
        else
            return () -> {
                ThreadLocalManager.setContext(previousContext);
                try {
                    return task.call();
                } finally {
                    ThreadLocalManager.clearContext();
                }
            };
    }
    public static Runnable wrapWithContext(Runnable task) {
        Map<String, Object> previousContext = ThreadLocalManager.getContext();
        if (previousContext == null) {
            return task;
        } else
            return () -> {
                ThreadLocalManager.setContext(previousContext);
                try {
                    task.run();
                } finally {
                    ThreadLocalManager.clearContext();
                }
            };
    }
}

Using the methods from ContextUtility we create CustomThreadPoolExecutor to override methods to attach thread context before submitting/executing tasks

CustomThreadPoolExecutor

public class CustomThreadPoolExecutor extends ThreadPoolExecutor {
    public CustomThreadPoolExecutor(int corePoolSize,
            int maximumPoolSize,
            long keepAliveTime,
            @NotNull TimeUnit unit,
            @NotNull BlockingQueue<Runnable> workQueue,
            @NotNull ThreadFactory threadFactory,
            @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }
    @Override
    public <T> @NotNull Future<T> submit(@NotNull Callable<T> task) {
        return super.submit(ContextUtility.wrapWithContext(task));
    }
    @Override
    public <T> @NotNull Future<T> submit(@NotNull Runnable task, T result) {
        return super.submit(ContextUtility.wrapWithContext(task), result);
    }
    @Override
    public @NotNull Future<?> submit(@NotNull Runnable task) {
        return super.submit(ContextUtility.wrapWithContext(task));
    }
    @Override
    public void execute(@NotNull Runnable task) {
        super.execute(ContextUtility.wrapWithContext(task));
    }

Now, we will use this executor in our custom MDC to allow creating custom dispatchers.

CustomDispatcherConfigurator

public class CustomDispatcherConfigurator extends MessageDispatcherConfigurator {
    private final CustomDispatcher instance;
    public CustomDispatcherConfigurator(Config config, DispatcherPrerequisites prerequisites) {
        super(config, prerequisites);
        Config threadPoolConfig = config.getConfig("thread-pool-executor");
        int fixedPoolSize = threadPoolConfig.getInt("fixed-pool-size");
        instance = new CustomDispatcher(
                this,
                config.getString("id"),
                config.getInt("throughput"),
                Duration.create(config.getDuration("throughput-deadline-time", TimeUnit.NANOSECONDS), TimeUnit.NANOSECONDS),
                (id, threadFactory) -> () -> new CustomThreadPoolExecutor(fixedPoolSize,
                                                                          fixedPoolSize,
                                                                          threadPoolConfig.getDuration("keep-alive-time", TimeUnit.MILLISECONDS),
                                                                          TimeUnit.MILLISECONDS,
                                                                          new LinkedBlockingDeque<>(),
                                                                          new ThreadFactory() {
                                                                              private int threadId = 1;
                                                                              @Override
                                                                              public Thread newThread(@NotNull Runnable r) {
                                                                                  Thread thread = new Thread(r);
                                                                                  thread.setName(config.getString("name") + "-" + threadId++);
                                                                                  return thread;
                                                                              }
                                                                          }),
                Duration.create(config.getDuration("shutdown-timeout", TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS)
        );
    }
    @Override
    public MessageDispatcher dispatcher() {
        return instance;
    }
}
class CustomDispatcher extends Dispatcher {
    public CustomDispatcher(MessageDispatcherConfigurator _configurator,
                            String id,
                            int throughput,
                            Duration throughputDeadlineTime,
                            ExecutorServiceFactoryProvider executorServiceFactoryProvider,
                            scala.concurrent.duration.FiniteDuration shutdownTimeout) {
        super(_configurator, id, throughput, throughputDeadlineTime, executorServiceFactoryProvider, shutdownTimeout);
    }
}

Now we can create different actors in our actor system using this custom MDC and define their config in aplication.conf

db-io-dispatcher {
    type = "contexts.CustomDispatcherConfigurator"
    executor = "thread-pool-executor"
    thread-pool-executor {
        fixed-pool-size = 11
    }
    throughput = 1
    shutdown-timeout = 60s
}
web-io-dispatcher {
    type = "contexts.CustomDispatcherConfigurator"
    executor = "thread-pool-executor"
    thread-pool-executor {
        fixed-pool-size = 20
    }
    throughput = 1
    shutdown-timeout = 60s
}

To create an actor for db execution context

DatabaseIODispatcher

public class DatabaseIODispatcher extends CustomExecutionContext {
    @Inject
    public DatabaseIODispatcher(ActorSystem actorSystem) {
        super(actorSystem, "db-io-dispatcher");
    }
}

This allows us to switch context from one executor to another without losing thread’s context.

Points to Remember

In async implementations, if the thread is switched in-between from custom executors to default ThreadPoolExecutor or ForkJoinPool, the thread context will get lost, hence we need to ensure that the thread context is not getting lost if any library method is using default pools.
We need to clear thread context after the task is complete or else it can cause memory leaks or OOM issues.

Thanks for reading. I hope this helps the community to provide more transparency in their applications.

Enhancing Transparency and Accountability: Implementing Entity Audit Logging in Java [Part-1]

Anuj Singh — Sat, 22 Jun 2024 09:58:10 +0000

In the dynamic world of Java development, I embarked on a mission to fulfil a client’s request for comprehensive user activity logs across our company’s services. To achieve this, I turned to Ebean’s changelog feature, a robust tool for audit logging in Java applications.

Initially, I experimented with manual logging by modifying the save method of our model entities. However, this approach proved cumbersome and prone to errors, lacking the comprehensive tracking capabilities we required. Additionally, I explored Ebean’s History annotation, hoping it would streamline the audit logging process. Unfortunately, it didn’t meet our needs as it lacked the granularity and customisation options necessary for our application.

Undeterred, I focused on leveraging Ebean’s changelog functionality to seamlessly integrate audit logging into our model entities. Each modification became a part of a detailed activity log, providing insight into users’ actions across our services.

As the development progressed, the application transformed into a testament to our commitment to data integrity. Stakeholders applauded the newfound visibility into user activity, while compliance requirements were effortlessly met. Through the power of Ebean’s Changelog and my dedication to meeting client needs, our Java application emerged as a beacon of transparency, paving the way for a more accountable future in our services.

BASIC ENTITY AUDIT LOGS

To enable audit logging for an entity we need to first enable @ChangeLog annotation from io.ebean.annotation on the entity. By default inserts are included. It can be excluded if required using inserts = ChangeLogInsertMode.EXCLUDE option in the annotation.

@Entity
@ChangeLog
@Table(name = "datagroup")
public class DataGroup extends Model {
  ...
@ChangeLog(inserts = ChangeLogInsertMode.EXCLUDE)
@Entity
public class Dataset extends Model {

If you want to log entity change only on certain field changes, then we can use

@ChangeLog(updatesThatInclude = {"field1","field2"})

Then we need to define a log appender for ChangeLogs in logback.xml. In our example we are using a ConsoleAppender to log out to STDOUT, however there are many options to log to socket, Files, DB, Kafka, SMTP or even slack and other apps. Refer this link for various appenders from apache and its not limited to this. There are other libraries offering more options as well.

<appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
    <encoder>
        <pattern>%date{yyyy-MM-dd HH:mm:ss} %coloredLevel %logger{15} - %message%n%xException{10}</pattern>
    </encoder>
</appender>
<logger name="io.ebean.ChangeLog" level="TRACE" additivity="false">
    <appender-ref ref="STDOUT"/>
</logger>

This is sufficient to log entity changes (inserts and updates).

AUDIT LOGS WITH USER CONTEXT

To log user context or additional info along with the changes to the entity, we need to implement ChangeLogPrepare class to add the additional info to the change set. The implementation of ChangeLogPrepare can be automatically detected if classpath scanning is on (just like entity beans are found). That is, if scanning is on we don’t need to explicitly register the ChangeLogPrepare implementation and instead it will be found and instantiated. We can also register it explicitly in DatabaseConfig as used in this doc.

For our use-case, we’ll keep the implementations inside models folder.

class MyChangeLogPrepare implements ChangeLogPrepare {
  @Override
  public boolean prepare(ChangeSet changes) {
// get user context information typically from a ThreadLocal or similar mechanism
    String currentUserId = ...;
    changes.setUserId(currentUserId);
    String userIpAddress = ...;
    changes.setUserIpAddress(userIpAddress);
    changes.setSource("myApplicationName");
// add arbitrary user context information to the userContext map
    changes.getUserContext().put("some", "thing");
    return true;
  }
}

In the next post, we'll take a look into how all this fits in an asynchronous application where execution contexts can change frequently and the change context might get lost in thread switching.