Skip to content

Commit a2db6ad

Browse files
hellojavaerhellojavaer
committed
DBClusterRouteContext and ShardRouteRuleExpressionContext support context stack
1 parent 1939f65 commit a2db6ad

File tree

3 files changed

+152
-36
lines changed

3 files changed

+152
-36
lines changed

‎ddal-ddr/src/main/java/org/hellojavaer/ddal/ddr/cluster/DBClusterRouteContext.java‎

Lines changed: 82 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import java.lang.reflect.Method;
2222
import java.util.HashMap;
23+
import java.util.LinkedList;
2324
import java.util.Map;
2425

2526
/**
@@ -28,68 +29,128 @@
2829
*/
2930
public class DBClusterRouteContext {
3031

31-
private static final ThreadLocal<String> clusterNameCache = new ThreadLocal();
32+
private static final Map<String, Object> systemVariables = new HashMap<>();
3233

33-
private static final Map<String, Object> systemVariables = new HashMap<>();
34+
private static final ThreadLocal<LinkedList<Context>> STACK = new ThreadLocal<LinkedList<Context>>() {
3435

35-
private static final ThreadLocal<Map<String, Object>> variables = new ThreadLocal() {
36-
37-
@Override
38-
protected Map initialValue() {
39-
return new HashMap();
40-
}
41-
};
36+
protected LinkedList<Context> initialValue() {
37+
LinkedList stack = new LinkedList<>();
38+
stack.add(new Context());
39+
return stack;
40+
}
41+
};
4242

4343
static {
4444
setSystemVariable(MathFunction.class);
4545
setSystemVariable(FormatFunction.class);
4646
}
4747

48+
public static void setClusterName(String name) {
49+
getCurrentContext().setClusterName(name);
50+
}
51+
52+
public static String getClusterName() {
53+
return getCurrentContext().getClusterName();
54+
}
55+
4856
private static void setSystemVariable(Class<?> clazz) {
4957
Method[] methods = clazz.getDeclaredMethods();
5058
for (Method method : methods) {
51-
systemVariables.put(method.getName(), method);
59+
setSystemVariable(method.getName(), method);
5260
}
5361
}
5462

55-
public static void setClusterName(String name) {
56-
clusterNameCache.set(name);
63+
private static void setSystemVariable(String name, Object value) {
64+
systemVariables.put(name, value);
5765
}
5866

59-
public static String getClusterName() {
60-
return clusterNameCache.get();
67+
public static Object getSystemVariable(String name) {
68+
return systemVariables.get(name);
6169
}
6270

6371
public static Object getVariable(String name) {
6472
if (name == null) {
6573
throw new IllegalArgumentException("name can't be null");
6674
}
67-
Object value = variables.get().get(name);
75+
for (Context context : STACK.get()) {
76+
Object value = context.getLocalVariables().get(name);
77+
if (value != null) {
78+
return value;
79+
}
80+
}// else
81+
return systemVariables.get(name);
82+
}
83+
84+
public static Object getLocalVariable(String name) {
85+
if (name == null) {
86+
throw new IllegalArgumentException("name can't be null");
87+
}
88+
Object value = getCurrentContext().getLocalVariables().get(name);
6889
if (value == null) {
6990
value = systemVariables.get(name);
7091
}
7192
return value;
7293
}
7394

74-
public static Object setVariable(String name, Object value) {
95+
public static Object setLocalVariable(String name, Object value) {
7596
if (name == null) {
7697
throw new IllegalArgumentException("name can't be null");
7798
}
7899
if (value == null) {
79100
throw new IllegalArgumentException("value can't be null");
80101
}
81-
return variables.get().put(name, value);
102+
return getCurrentContext().getLocalVariables().put(name, value);
82103
}
83104

84-
public static Object removeVariable(String name) {
105+
public static Object removeLocalVariable(String name) {
85106
if (name == null) {
86107
throw new IllegalArgumentException("name can't be null");
87108
}
88-
return variables.get().remove(name);
109+
return getCurrentContext().getLocalVariables().remove(name);
110+
}
111+
112+
public static void pushContext() {
113+
STACK.get().addFirst(new Context());
114+
}
115+
116+
public static void popContext() throws IndexOutOfBoundsException {
117+
if (STACK.get().size() <= 1) {
118+
throw new IndexOutOfBoundsException("root context can't be pop");
119+
} else {
120+
STACK.get().removeFirst();
121+
}
89122
}
90123

91124
public static void clearContext() {
92-
clusterNameCache.remove();
93-
variables.get().clear();
125+
Context context = getCurrentContext();
126+
context.getLocalVariables().clear();
127+
context.setClusterName(null);
128+
}
129+
130+
private static Context getCurrentContext() {
131+
return STACK.get().getFirst();
132+
}
133+
134+
private static class Context {
135+
136+
private String clusterName = null;
137+
private Map<String, Object> localVariables = new HashMap<>();
138+
139+
public String getClusterName() {
140+
return clusterName;
141+
}
142+
143+
public void setClusterName(String clusterName) {
144+
this.clusterName = clusterName;
145+
}
146+
147+
public Map<String, Object> getLocalVariables() {
148+
return localVariables;
149+
}
150+
151+
public void setLocalVariables(Map<String, Object> localVariables) {
152+
this.localVariables = localVariables;
153+
}
94154
}
155+
95156
}

‎ddal-ddr/src/main/java/org/hellojavaer/ddal/ddr/expression/ShardRouteRuleExpressionContext.java‎

Lines changed: 68 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import java.lang.reflect.Method;
2222
import java.util.HashMap;
23+
import java.util.LinkedList;
2324
import java.util.Map;
2425

2526
/**
@@ -29,55 +30,108 @@
2930
public class ShardRouteRuleExpressionContext {
3031

3132
private static final Map<String, Object> systemVariables = new HashMap<>();
32-
private static final ThreadLocal<Map<String, Object>> variables = new ThreadLocal() {
3333

34-
@Override
35-
protected Map initialValue() {
36-
return new HashMap();
34+
private static final ThreadLocal<LinkedList<Context>> STACK = new ThreadLocal<LinkedList<Context>>() {
35+
36+
protected LinkedList<Context> initialValue() {
37+
LinkedList stack = new LinkedList<>();
38+
stack.add(new Context());
39+
return stack;
3740
}
3841
};
3942

4043
static {
41-
setMethodVariable(MathFunction.class);
42-
setMethodVariable(FormatFunction.class);
44+
setSystemVariable(MathFunction.class);
45+
setSystemVariable(FormatFunction.class);
4346
}
4447

45-
private static void setMethodVariable(Class<?> clazz) {
48+
private static void setSystemVariable(Class<?> clazz) {
4649
Method[] methods = clazz.getDeclaredMethods();
4750
for (Method method : methods) {
48-
systemVariables.put(method.getName(), method);
51+
setSystemVariable(method.getName(), method);
4952
}
5053
}
5154

55+
private static void setSystemVariable(String name, Object value) {
56+
systemVariables.put(name, value);
57+
}
58+
59+
public static Object getSystemVariable(String name) {
60+
return systemVariables.get(name);
61+
}
62+
5263
public static Object getVariable(String name) {
5364
if (name == null) {
5465
throw new IllegalArgumentException("name can't be null");
5566
}
56-
Object value = variables.get().get(name);
67+
for (Context context : STACK.get()) {
68+
Object value = context.getLocalVariables().get(name);
69+
if (value != null) {
70+
return value;
71+
}
72+
}// else
73+
return systemVariables.get(name);
74+
}
75+
76+
public static Object getLocalVariable(String name) {
77+
if (name == null) {
78+
throw new IllegalArgumentException("name can't be null");
79+
}
80+
Object value = getCurrentContext().getLocalVariables().get(name);
5781
if (value == null) {
5882
value = systemVariables.get(name);
5983
}
6084
return value;
6185
}
6286

63-
public static Object setVariable(String name, Object value) {
87+
public static Object setLocalVariable(String name, Object value) {
6488
if (name == null) {
6589
throw new IllegalArgumentException("name can't be null");
6690
}
6791
if (value == null) {
6892
throw new IllegalArgumentException("value can't be null");
6993
}
70-
return variables.get().put(name, value);
94+
return getCurrentContext().getLocalVariables().put(name, value);
7195
}
7296

73-
public static Object removeVariable(String name) {
97+
public static Object removeLocalVariable(String name) {
7498
if (name == null) {
7599
throw new IllegalArgumentException("name can't be null");
76100
}
77-
return variables.get().remove(name);
101+
return getCurrentContext().getLocalVariables().remove(name);
102+
}
103+
104+
public static void pushContext() {
105+
STACK.get().addFirst(new Context());
106+
}
107+
108+
public static void popContext() throws IndexOutOfBoundsException {
109+
if (STACK.get().size() <= 1) {
110+
throw new IndexOutOfBoundsException("root context can't be pop");
111+
} else {
112+
STACK.get().removeFirst();
113+
}
78114
}
79115

80116
public static void clearContext() {
81-
variables.get().clear();
117+
Context context = getCurrentContext();
118+
context.getLocalVariables().clear();
119+
}
120+
121+
private static Context getCurrentContext() {
122+
return STACK.get().getFirst();
123+
}
124+
125+
private static class Context {
126+
127+
private Map<String, Object> localVariables = new HashMap<>();
128+
129+
public Map<String, Object> getLocalVariables() {
130+
return localVariables;
131+
}
132+
133+
public void setLocalVariables(Map<String, Object> localVariables) {
134+
this.localVariables = localVariables;
135+
}
82136
}
83137
}

‎ddal-spring/src/main/java/org/hellojavaer/ddal/spring/scan/EnableDBClusterRouteAnnotation.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ public class EnableDBClusterRouteAnnotation {
4545
@Around("@annotation(dbClusterRoute)")
4646
public Object around(ProceedingJoinPoint joinPoint, DBClusterRoute dbClusterRoute) throws Throwable {
4747
try {
48+
DBClusterRouteContext.pushContext();
4849
Object[] args = joinPoint.getArgs();
4950
MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
5051
Method method = methodSignature.getMethod();
@@ -62,7 +63,7 @@ public Object around(ProceedingJoinPoint joinPoint, DBClusterRoute dbClusterRout
6263
DBClusterRouteContext.setClusterName(targetClusterName);
6364
return joinPoint.proceed(args);
6465
} finally {
65-
DBClusterRouteContext.clearContext();
66+
DBClusterRouteContext.popContext();
6667
}
6768
}
6869
}

0 commit comments

Comments
 (0)