From bcfdd1c54f7e265e43691325eb99eda554a2d4db Mon Sep 17 00:00:00 2001
From: Filip Hrisafov <filip.hrisafov@gmail.com>
Date: Mon, 9 Oct 2023 15:46:51 +0200
Subject: [PATCH] Use a single thread local with a custom object for the tenant

---
 .../impl/tenant/ThreadLocalTenantContext.java | 19 +++++++++----------
 .../tenant/ThreadLocalTenantContextTest.java  | 13 +++++--------
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/modules/flowable-engine-common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java b/modules/flowable-engine-common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java
index 347b8b893b2..9d9e08d5e3a 100644
--- a/modules/flowable-engine-common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java
+++ b/modules/flowable-engine-common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java
@@ -12,7 +12,6 @@
  */
 package org.flowable.common.engine.impl.tenant;
 
-import org.flowable.common.engine.api.FlowableException;
 import org.flowable.common.engine.api.tenant.TenantContext;
 
 /**
@@ -20,30 +19,30 @@
  */
 public class ThreadLocalTenantContext implements TenantContext {
 
-    protected final ThreadLocal<String> tenantId = ThreadLocal.withInitial(() -> {
-        throw new FlowableException("Tenant value has not been set");
-    });
-    protected final ThreadLocal<Boolean> tenantIdSet = ThreadLocal.withInitial(() -> Boolean.FALSE);
+    protected final ThreadLocal<Tenant> tenantId = new ThreadLocal<>();
 
     @Override
     public String getTenantId() {
-        return tenantId.get();
+        Tenant tenant = tenantId.get();
+        return tenant != null ? tenant.tenantId() : null;
     }
 
     @Override
     public void setTenantId(String tenantId) {
-        this.tenantId.set(tenantId);
-        this.tenantIdSet.set(true);
+        this.tenantId.set(new Tenant(tenantId));
     }
 
     @Override
     public void clearTenantId() {
         tenantId.remove();
-        tenantIdSet.remove();
     }
 
     @Override
     public boolean isTenantIdSet() {
-        return tenantIdSet.get();
+        return tenantId.get() != null;
+    }
+
+    protected record Tenant(String tenantId) {
+
     }
 }
diff --git a/modules/flowable-engine-common/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java b/modules/flowable-engine-common/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java
index 1612393eedf..9f149ab8811 100644
--- a/modules/flowable-engine-common/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java
+++ b/modules/flowable-engine-common/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java
@@ -13,9 +13,7 @@
 package org.flowable.common.engine.impl.tenant;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
-import org.flowable.common.engine.api.FlowableException;
 import org.flowable.common.engine.api.tenant.TenantContext;
 import org.junit.jupiter.api.Test;
 
@@ -28,9 +26,7 @@ class ThreadLocalTenantContextTest {
 
     @Test
     void getTenantId() {
-        assertThatThrownBy(() -> underTest.getTenantId())
-                .isExactlyInstanceOf(FlowableException.class)
-                .hasMessage("Tenant value has not been set");
+        assertThat(underTest.getTenantId()).isNull();
 
         underTest.setTenantId("acme");
         assertThat(underTest.getTenantId()).isEqualTo("acme");
@@ -39,9 +35,7 @@ void getTenantId() {
         assertThat(underTest.getTenantId()).isEqualTo("muppets");
 
         underTest.clearTenantId();
-        assertThatThrownBy(() -> underTest.getTenantId())
-                .isExactlyInstanceOf(FlowableException.class)
-                .hasMessage("Tenant value has not been set");
+        assertThat(underTest.getTenantId()).isNull();
     }
 
     @Test
@@ -53,5 +47,8 @@ void isTenantIdSet() {
 
         underTest.clearTenantId();
         assertThat(underTest.isTenantIdSet()).isFalse();
+
+        underTest.setTenantId(null);
+        assertThat(underTest.isTenantIdSet()).isTrue();
     }
 }