Spring 代理搭建分布式锁组件

Spring 代理搭建分布式锁组件

实习的时候在公司看代码，看到了一个使用 Spring 的代理和 SPI 机制等做的一个分布式锁组件，觉得很有意思

该组件的主要作用就是通过注解@DistributedLock来实现分布式锁，并使用SPI来自定义分布式锁的具体逻辑实现

简单实现

这里写一个简单的Demo，首先先定义注解

/**
 * 分布式锁注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Inherited
public @interface DistributedLock {
    String value() default "";

    /**
     * 分布式锁key
     */
    @AliasFor("value")
    String lockKey() default "";
}

该注解起作用是通过 Spring 的BeanPostProcessor，扫描 Bean 是否标有该注解，有则生成代理对象

/**
 * 对标记有 @DistributedLock 注解的Bean生成对应的代理对象
 */
public class DistributedLockPostProcessor implements BeanPostProcessor {
    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        Method[] methods = bean.getClass().getMethods();
        List<Method> methodList = Arrays.asList(methods).stream().filter(method -> {
            DistributedLock distributeLock = method.getAnnotation(DistributedLock.class);
            return Objects.nonNull(distributeLock);
        }).toList();
        // 没有标记注解则直接返回
        if (methodList.isEmpty()) {
            return bean;
        }
        DistributedLockInterceptor distributedLockInterceptor = loadDistributedLockInterceptor();
        return buildProxy(bean, distributedLockInterceptor);
    }

    /**
     * 加载 Interceptor，用于构建代理对象
     */
    private DistributedLockInterceptor loadDistributedLockInterceptor() {
        // 使用 SPI 机制获取 Interceptor 对象
        ServiceLoader<DistributedLockInterceptor> distributedLockInterceptors = ServiceLoader.load(DistributedLockInterceptor.class);
        List<DistributedLockInterceptor> distributedLockInterceptorList = distributedLockInterceptors.stream().map(ServiceLoader.Provider::get).toList();
        return distributedLockInterceptorList.get(0);
    }

    /**
     * 构建代理对象
     */
    private Object buildProxy(Object bean, DistributedLockInterceptor interceptor) {
        ProxyFactory proxyFactory = new ProxyFactory(bean);
        // 启用CGLIB代理
        proxyFactory.setProxyTargetClass(true);
        proxyFactory.addAdvice(interceptor);
        return proxyFactory.getProxy(bean.getClass().getClassLoader());
    }
}

这里在生成代理对象是有一个坑，就是要通过proxyFactory.setProxyTargetClass(true);启用 CGLIB 代理，否则在生成代理时，由于 Service 接口没有标记注解，注解在实现类上，从而导致拿不到注解信息

通过ServiceLoader.load()使用SPI加载实现类，需要在resources.META-INF下编写配置，这样用户就能通过该配置自定义分布式锁的具体实现，从而达到解耦和自定义配置的目的

里面写下具体实现类的全类名

cn.bobasyu.distribute.proxy.DefaultDoNothingDistributedLockInterceptor

DistributedLockInterceptor是一个抽象类，只定义了代理对象中加锁解锁的基本逻辑，下面是简单的代码

public abstract class DistributedLockInterceptor implements MethodInterceptor {

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        Method method = invocation.getMethod();
        DistributedLock distributedLock = method.getAnnotation(DistributedLock.class);
        if (Objects.isNull(distributedLock)) {
            return invocation.proceed();
        }
        Object argument = invocation.getArguments()[0];
        Object lockKey = parseLockKey(argument, distributedLock.lockKey());
        try {
            buildLock(lockKey);
            boolean isLocked = tryLock();
            if (!isLocked) {
                throw new RuntimeException("加锁失败，请稍后再试");
            }
            return invocation.proceed();
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            unLock();
        }
    }

    /**
     * 解析分布式锁key
     */
    private Object parseLockKey(Object argumentObject, String lockKeyExpression) {
        ExpressionParser parser = new SpelExpressionParser();
        StandardEvaluationContext context = new StandardEvaluationContext(argumentObject);
        Expression expression = parser.parseExpression(lockKeyExpression);
        return expression.getValue(context);
    }

    protected abstract void buildLock(Object lockKey);

    protected abstract boolean tryLock();

    protected abstract void unLock();
}

在parseLockKey方法中，使用了 SpEL 的方式解析来获取 key，这里的逻辑是解析调用代理方法的第一个参数，调用其指定方法，其返回值为分布式锁的 key

在这个抽象类中，定义了buildLock，tryLock，unLock三个抽象类，其交给子类来具体实现，抽象类中只定义了基本逻辑，下面是其一个实现类DefaultDoNothingDistributedLockInterceptor，除了输出信息没有做其他任何动作

public class DefaultDoNothingDistributedLockInterceptor extends DistributedLockInterceptor {
    private Object lockKey;

    @Override
    protected void buildLock(Object lockKey) {
        System.out.println("获得锁，lockKey=" + lockKey);
        this.lockKey = lockKey;

    }

    @Override
    protected boolean tryLock() {
        System.out.println("尝试加锁，lockKey=" + lockKey);
        return true;
    }

    @Override
    protected void unLock() {
        System.out.println("解锁，lockKey=" + lockKey);
    }
}

然后是基本配置类，这里只需要加载DistributedLockPostProcessor到容器中即可

@Configuration
public class DistributedLockConfiguration {
    @Bean
    public DistributedLockPostProcessor distributedLockPostProcessor() {
        return new DistributedLockPostProcessor();
    }
}

也可以接着使用spring.factories进行配置，这样导入组件时就自动注入了，或者在使用者的配置类中使用@ComponentScan注解等方式主动扫描

使用测试

具体使用，先定义Service：

@Service
public class TestServiceImpl implements TestService {

    @Override
    @DistributedLock(lockKey = "getUserType().getTypeName()")
    public void doSomething(UserInfo userInfo) {
        System.out.println("do something.");
    }
}

前面说到，其使用了SpEL的方式解析获得分布式锁key，在这里就是调用userInfo.getUserType().getTypeName()，其返回值作为key

然后是测试：

@SpringBootTest
class SpringDistributedApplicationTests {
   @Autowired
   TestService testService;

   @Test
   void contextLoads() {
       UserInfo userInfo = new UserInfo("AAA", UserType.USER_TYPE1);
       testService.doSomething(userInfo);
   }
}

输出：

获得锁，lockKey=userType1
尝试加锁，lockKey=userType1
do something.
解锁，lockKey=userType1

个人想法：

整体看下来相当于做了个AOP，调用前加锁，调用结束后解锁，这里完全可以直接使用 Spring AOP 直接实现，使用环绕通知，直接对标有@DistributedLock注解的地方进行切入，这样就需要编写PostProcessor和 MethodInteceptor，SPI 也可以直接让用户使用 Configuration 主动配置