详解Reactor中Context的用法

在响应式编程中，多线程异步性成为天然的内在，多线程之间的切换也成为原生的，在处理一个数据流Flux/Mono时，基本无法知道是运行在哪个线程上或哪个线程池里，可以说，每一个操作符operator以及内部的函数都可能运行在不同的线程上。这就意味着，以前用ThreadLocal来作为方法间透明传递共享变量的方式不再行得通。为此，Reactor提供了Context来替代ThreadLocal实现一个跨线程的共享变量的透明方式。

本文会从以下几个方面来介绍Context的相关知识：

一、使用介绍

static String KEY = "TEST_CONTEXT_KEY";
static String KEY2 = "TEST_CONTEXT_KEY2";

public static void main(String[] args) {
	Flux flux = convert("hello", Flux.just(1, 2, 3));
	flux
    	.subscriberContext(Context.of(KEY, "Outside"))
    	.subscribe(v -> System.out.println(v));
}

public static Flux convert(String prefix, Flux publisher) {
    return publisher.map(v -> prefix + " " + v)
        .subscriberContext(Context.of(KEY, "NotUsed"))
        .flatMap(v -> Mono.subscriberContext().map(ctx -> ctx.get(KEY) + " " + ctx.get(KEY2) + " " + v))
        .subscriberContext(context -> context.put(KEY2, "Inside"))
        .flatMap(v -> Mono.subscriberContext().map(ctx -> ctx.get(KEY) + " " + v));
}

上面是context的使用方案介绍，其输出如下：

Outside Outside Inside hello 1
Outside Outside Inside hello 2
Outside Outside Inside hello 3

上面的使用案例展示了一个使用context的常见例子。通过在外部方法里传入context，如flux.subscriberContext(Context.of(KEY, "Outside"))，使得内部方法convert能够获取外界环境的context，同时内部方法还可以增加自己的context数据，如subscriberContext(context -> context.put(KEY2, "Inside"))，结合之后，在让内部的方法（flatMap里的方法）感知到整个上下文context的数据内容。

对于context的使用，主要分为几个部分： 1. context的创建 2. context的写入（传入）与读取 3. 执行顺序

1. context —— 不可变对象

由于reactor天然是跨线程的，所以context设计为了不可变的对象，即每次的更新都是创建一个新的对象。每次的put/putAll操作，都是先把旧对象的值复制到新对象，然后再进行put/putAll等更新操作。

2. context的写入与读取

context写入是使用subscriberContext方法，其入参有两种形式：传值方式subscriberContext(ctx)与lambda函数方式 —— subscriberContext(ctx -> ctx.put(key,value))。

context的读取是利用Mono的静态方法subscriberContext()来获取，由于其返回的是一个Mono, 所以通常与flatMap结合使用。

3. 执行顺序

context的传入是发生在subscribe()订阅阶段的，所以其写入的顺序是从下往上的，即在示例中，先执行subscriberContext(Context.of(KEY, "Outside"))，再执行subscriberContext(context -> context.put(KEY2, "Inside")), 最后执行subscriberContext(Context.of(KEY, "NotUsed"))。在订阅阶段执行完后，进入运行阶段，数据流从上往下执行，每次读取context的时候Mono.subscriberContext()都是读取下一个的context。所以"NotUsed"的context并没有生效。

此外，context.put()操作是复制旧的再update新的对象，所以Mono.subscriberContext().map(ctx -> ctx.get(KEY) + " " + ctx.get(KEY2) + " " + v)这个阶段仍能读取前一个context关于KEY的内容。

总结

注意

subscriberContext(Context.of("Outside")与subscriberContext(context -> Context.of("Outside"))是有区别，前者是会结合复用前面的context，而后者是直接返回一个新的context并不会复用前面的context。 其原因是，subscriberContext(Context.of("Outside")) 其实内部调用的是subscriberContext(context -> context.putAll(Context.of("Outside"))，其入参的context就是前面的context，putAll方法会复用前面的context。而 subscriberContext(context -> Context.of("Outside"))不复用的原因就是因为放弃了入参的context。所以，可以利用这种方式来放弃之前的context，当然不鼓励这么做，因为你不清楚之前context会不会影响后续的程序。

本文章的代码用的事reactor 3.3的版本，自3.5之后，subscriberContext方法改为contextWrite，读取的方法改为deferContextual。

二、源码解读

现在我们从源代码上看看，context写入为什么是自下而上的，读取的时候又是依附于下一个subscriber并且自上而下的。

public final Flux subscriberContext(Function doOnContext) {
	return new FluxContextStart<>(this, doOnContext);
}
FluxContextStart(Flux source, Function doOnContext) {
	super(source);
	this.doOnContext = Objects.requireNonNull(doOnContext, "doOnContext");
}

@Override
public CoreSubscriber subscribeOrReturn(CoreSubscriber actual) {
    Context c = doOnContext.apply(actual.currentContext());
    return new ContextStartSubscriber<>(actual, c);
}
ContextStartSubscriber(CoreSubscriber actual, Context context) {
    this.actual = actual;
    this.context = context;
    if (actual instanceof ConditionalSubscriber) {
        this.actualConditional = (ConditionalSubscriber) actual;
    }
    else {
        this.actualConditional = null;
    }
}
@Override
public Context currentContext() {
    return this.context;
}

上面截取了subscriberContext方法的源代码，可以看到subscriberContext方法最终会创建ContextStartSubscriber的对象，并将生成的context赋值Context c = doOnContext.apply(actual.currentContext())，所以context是伴随subscriberContext方法对应的subscriber里的。

由于context赋值操作Context c = doOnContext.apply(actual.currentContext())是发生在subscribeOrReturn方法里，即发生在subscribe()订阅阶段，所以整个执行的顺序是自下而上的（沿着整个flow自下而上至源头的publisher）。

那读取context的时候为什么是自上而下的呢？我们来看下读取操作Mono.subscribeContext()的源码。

public static Mono subscriberContext() {
    return onAssembly(MonoCurrentContext.INSTANCE);
}
final class MonoCurrentContext extends Mono
		implements Fuseable, Scannable {
	static final MonoCurrentContext INSTANCE = new MonoCurrentContext();
	public void subscribe(CoreSubscriber actual) {
		Context ctx = actual.currentContext();
		actual.onSubscribe(Operators.scalarSubscription(actual, ctx));
	}
}
interface InnerOperator
		extends InnerConsumer, InnerProducer {
	@Override
	default Context currentContext() {
		return actual().currentContext();
	}
}

logging.pattern.level: "[%X{X-Request-Id}] [%thread] %-5level - %msg %n"

@SpringBootApplication
@Slf4j
@RestController
public class MdcApplication {

  public static void main(String[] args) {
    SpringApplication.run(MdcApplication.class, args);
  }

  private final static String X_REQUEST_ID_KEY = "X-Request-Id";

  @GetMapping("/")
  Flux split(@RequestParam("value") String value, @RequestHeader(X_REQUEST_ID_KEY) String requestId) {
    return Flux.fromArray(value.split("_"))
        .doOnEach(logWithContext(ch -> log.info("handling one item: {}", ch)))
        .subscriberContext(Context.of(X_REQUEST_ID_KEY, requestId));
  }

  private static  Consumer> logWithContext(Consumer logStatement) {
    return signal -> {
      if (!signal.isOnNext()) {
        return;
      }
      String requestId = signal.getContext().get(X_REQUEST_ID_KEY);
      try (MDC.MDCCloseable closeable = MDC.putCloseable(X_REQUEST_ID_KEY, requestId)) {
        logStatement.accept(signal.get());
      }
    };
  }
}

try {
	MDC.put(X_REQUEST_ID_KEY, requestId);
	logStatement.accept(signal.get());
} finally {
	MDC.remove(X_REQUEST_ID_KEY);
}

Mono.subscriberContext().map(ctx -> {   
    RequestHeadersSpec request = webClient.get().uri(uri);       
    request = request.header("X-Request-ID", ctx.get(X_REQUEST_ID_KEY)); 
    // The rest of your request logic...    
});