Java Code Examples for com.linecorp.armeria.server.ServiceRequestContext#setAttr()

private static ServiceRequestContext newServiceContext(
        String path, @Nullable String query) throws Exception {

    final InetSocketAddress remoteAddress = new InetSocketAddress(
            InetAddress.getByAddress("client.com", new byte[] { 1, 2, 3, 4 }), 5678);
    final InetSocketAddress localAddress = new InetSocketAddress(
            InetAddress.getByAddress("server.com", new byte[] { 5, 6, 7, 8 }), 8080);

    final String pathAndQuery = path + (query != null ? '?' + query : "");
    final HttpRequest req = HttpRequest.of(RequestHeaders.of(HttpMethod.GET, pathAndQuery,
                                                             HttpHeaderNames.AUTHORITY, "server.com:8080",
                                                             HttpHeaderNames.USER_AGENT, "some-client"));

    final ServiceRequestContext ctx =
            ServiceRequestContext.builder(req)
                                 .sslSession(newSslSession())
                                 .remoteAddress(remoteAddress)
                                 .localAddress(localAddress)
                                 .proxiedAddresses(
                                         ProxiedAddresses.of(new InetSocketAddress("9.10.11.12", 0)))
                                 .build();

    ctx.setAttr(MY_ATTR, new CustomObject("some-name", "some-value"));
    return ctx;
}
 

@Test
void canBringAttributeInServiceRequestContext() {
    final HttpRequest req = HttpRequest.of(HttpMethod.GET, "/");
    final ServiceRequestContext serviceContext = ServiceRequestContext.of(req);
    final AttributeKey<String> fooKey = AttributeKey.valueOf(DefaultClientRequestContextTest.class, "foo");
    serviceContext.setAttr(fooKey, "foo");
    try (SafeCloseable ignored = serviceContext.push()) {
        final ClientRequestContext clientContext = ClientRequestContext.of(req);
        assertThat(clientContext.attr(fooKey)).isEqualTo("foo");
        assertThat(clientContext.attrs().hasNext()).isTrue();

        final ClientRequestContext derivedContext = clientContext.newDerivedContext(
                clientContext.id(), clientContext.request(),
                clientContext.rpcRequest());
        assertThat(derivedContext.attr(fooKey)).isNotNull();
        // Attributes in serviceContext is not copied to clientContext when derived.

        final AttributeKey<String> barKey = AttributeKey.valueOf(DefaultClientRequestContextTest.class,
                                                                 "bar");
        clientContext.setAttr(barKey, "bar");
        assertThat(serviceContext.attr(barKey)).isNull();
    }
}
 

@Test
void shouldRepopulateWhenAttributeChanges() {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    final RequestContextExporter exporter =
            RequestContextExporter.builder()
                                  .addAttribute("attrs.attr1", ATTR1)
                                  .build();

    assertThat(exporter.export(ctx)).doesNotContainKeys("attrs.attr1");

    ctx.setAttr(ATTR1, "foo");
    assertThat(exporter.export(ctx)).containsEntry("attrs.attr1", "foo");

    ctx.setAttr(ATTR1, "bar");
    assertThat(exporter.export(ctx)).containsEntry("attrs.attr1", "bar");

    ctx.setAttr(ATTR1, null);
    assertThat(exporter.export(ctx)).doesNotContainKeys("attrs.attr1");
}
 

@Test
void shouldExportDifferentAliasOnSameKey() {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    ctx.setAttr(ATTR1, "1");
    ctx.setAttr(ATTR2, "2");
    final RequestContextExporter exporter =
            RequestContextExporter.builder()
                                  .addAttribute("attrs.attr1-1", ATTR1)
                                  .addAttribute("attrs.attr1-2", ATTR1)
                                  .addAttribute("attrs.attr2", ATTR2)
                                  .build();

    assertThat(exporter.export(ctx)).containsOnlyKeys(
            "attrs.attr1-1",
            "attrs.attr1-2",
            "attrs.attr2");
}
 

@Test
void customExportKey() {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    ctx.setAttr(ATTR1, "1");
    ctx.setAttr(ATTR3, new Foo("foo"));
    final RequestContextExporter exporter = RequestContextExporter
            .builder()
            .addAttribute("attrs.attr1", ATTR1)
            .addAttribute("my_attr2", ATTR1)
            .addRequestHeader(HttpHeaderNames.METHOD, "request_method")
            .addKeyPattern("request_id=req.id")
            .addKeyPattern("foo=attr:" + Foo.class.getName() + "#ATTR3")
            .addKeyPattern("bar=attr:" + Foo.class.getName() + "#ATTR3:" + FooStringifier.class.getName())
            .build();
    final Map<String, String> export;
    try (SafeCloseable ignored = ctx.push()) {
        export = exporter.export();
    }
    assertThat(export).containsOnlyKeys("request_id", "request_method",
                                        "attrs.attr1", "my_attr2",
                                        "foo", "bar");
}
 

@Test
void attrsDoNotIterateRootWhenKeyIsSame() {
    final HttpRequest req = HttpRequest.of(HttpMethod.GET, "/");
    final ServiceRequestContext serviceContext = ServiceRequestContext.of(req);
    try (SafeCloseable ignored = serviceContext.push()) {
        final ClientRequestContext clientContext = ClientRequestContext.of(req);
        final AttributeKey<String> fooKey = AttributeKey.valueOf(DefaultClientRequestContextTest.class,
                                                                 "foo");
        clientContext.setAttr(fooKey, "foo");
        serviceContext.setAttr(fooKey, "bar");
        final Iterator<Entry<AttributeKey<?>, Object>> attrs = clientContext.attrs();
        assertThat(attrs.next().getValue()).isEqualTo("foo");
        assertThat(attrs.hasNext()).isFalse();
    }
}
 

@Test
void shouldNotExportNullValue() {
    final ServiceRequestContext ctx = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    ctx.setAttr(ATTR1, "1");
    ctx.setAttr(ATTR2, null);
    ctx.logBuilder().endRequest();
    ctx.logBuilder().endResponse();
    final RequestContextExporter exporter =
            RequestContextExporter.builder()
                                  .addKeyPattern("*")
                                  .addAttribute("attrs.attr1", ATTR1)
                                  .addAttribute("attrs.attr2", ATTR2)
                                  .build();

    assertThat(exporter.export(ctx)).containsOnlyKeys(
            BuiltInProperty.CLIENT_IP.key,
            BuiltInProperty.ELAPSED_NANOS.key,
            BuiltInProperty.LOCAL_HOST.key,
            BuiltInProperty.LOCAL_IP.key,
            BuiltInProperty.LOCAL_PORT.key,
            BuiltInProperty.REMOTE_HOST.key,
            BuiltInProperty.REMOTE_IP.key,
            BuiltInProperty.REMOTE_PORT.key,
            BuiltInProperty.REQ_NAME.key,
            BuiltInProperty.REQ_SERVICE_NAME.key,
            BuiltInProperty.REQ_AUTHORITY.key,
            BuiltInProperty.REQ_CONTENT_LENGTH.key,
            BuiltInProperty.REQ_DIRECTION.key,
            BuiltInProperty.REQ_METHOD.key,
            BuiltInProperty.REQ_PATH.key,
            BuiltInProperty.RES_CONTENT_LENGTH.key,
            BuiltInProperty.RES_STATUS_CODE.key,
            BuiltInProperty.REQ_ID.key,
            BuiltInProperty.SCHEME.key,
            "attrs.attr1");
}
 

@Override
public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exception {
    ctx.setAttr(METADATA_SERVICE_ATTRIBUTE_KEY, mds);
    return unwrap().serve(ctx, req);
}
 

public static void setCurrentUser(ServiceRequestContext ctx, User currentUser) {
    requireNonNull(ctx, "ctx");
    requireNonNull(currentUser, "currentUser");
    ctx.setAttr(CURRENT_USER_KEY, currentUser);
}
 

private static void appendAttribute(ServiceRequestContext ctx, String value) {
    final String v = ctx.attr(decoratedFlag);
    ctx.setAttr(decoratedFlag, (v == null ? "" : v) + value);
}
 

@Test
void hasAttributeInRoot() {
    final ServiceRequestContext root = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    final DefaultAttributeMap child = new DefaultAttributeMap(root);

    // root: [foo], child: []
    final AttributeKey<String> foo = AttributeKey.valueOf("foo");
    root.setAttr(foo, "foo");

    assertThat(root.attr(foo)).isEqualTo("foo");
    assertThat(child.attr(foo)).isEqualTo("foo");
    assertThat(child.ownAttr(foo)).isNull();

    // root: [foo], child: [foo2]
    child.setAttr(foo, "foo2");
    assertThat(child.ownAttr(foo)).isEqualTo("foo2");
    assertThat(child.attr(foo)).isEqualTo("foo2");
    assertThat(root.attr(foo)).isEqualTo("foo");

    final AttributeKey<String> bar = AttributeKey.valueOf("bar");
    // root: [foo, bar], child: [foo2]
    root.setAttr(bar, "bar");
    // root: [foo, bar], child: [foo2, bar2]
    assertThat(child.setAttrIfAbsent(bar, "bar2")).isNull();

    assertThat(child.attr(bar)).isEqualTo("bar2");
    assertThat(root.attr(bar)).isEqualTo("bar");

    // Do not change.
    // root: [foo, bar], child: [foo2, bar2]
    assertThat(child.setAttrIfAbsent(bar, "bar3")).isEqualTo("bar2");
    assertThat(child.attr(bar)).isEqualTo("bar2");

    final AttributeKey<String> baz = AttributeKey.valueOf("baz");
    // root: [foo, bar, baz], child: [foo2, bar2]
    root.setAttr(baz, "baz");
    // root: [foo, bar, baz], child: [foo2, bar2, baz2]
    assertThat(child.computeAttrIfAbsent(baz, key -> "baz2")).isEqualTo("baz2");

    assertThat(child.attr(baz)).isEqualTo("baz2");
    assertThat(root.attr(baz)).isEqualTo("baz");

    // Do not change.
    // root: [foo, bar, baz], child: [foo2, bar2, baz2]
    assertThat(child.computeAttrIfAbsent(baz, key -> "baz3")).isEqualTo("baz2");
    assertThat(child.attr(baz)).isEqualTo("baz2");
}
 

@Test
void attrsWithRoot() {
    final ServiceRequestContext root = ServiceRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
    final DefaultAttributeMap child = new DefaultAttributeMap(root);

    final AttributeKey<String> foo = AttributeKey.valueOf("foo");
    // root: [foo], child: []
    root.setAttr(foo, "foo");

    Iterator<Entry<AttributeKey<?>, Object>> childIt = child.attrs();
    final Entry<AttributeKey<?>, Object> next = childIt.next();
    assertThat(next.getValue()).isEqualTo("foo");
    assertThat(childIt.hasNext()).isFalse();
    assertThatThrownBy(childIt::next).isInstanceOf(NoSuchElementException.class);

    // root: [foo], child: [foo1]
    next.setValue("foo1");
    assertThat(child.attr(foo)).isEqualTo("foo1");
    assertThat(child.ownAttr(foo)).isEqualTo("foo1");
    // The value of entry is changed.
    assertThat(next.getValue()).isEqualTo("foo1");
    // root attribute remains unaffected.
    assertThat(root.attr(foo)).isEqualTo("foo");

    // Set a new attribute to child.
    final AttributeKey<String> bar = AttributeKey.valueOf("bar");
    // root: [foo], child: [foo1, bar]
    child.setAttr(bar, "bar");

    childIt = child.attrs();
    final List<String> attributeValues = new ArrayList<>(2);
    Entry<AttributeKey<?>, Object> barEntry = null;
    for (int i = 0; i < 2; i++) {
        final Entry<AttributeKey<?>, Object> tempEntry = childIt.next();
        attributeValues.add((String) tempEntry.getValue());
        if ("bar".equals(tempEntry.getValue())) {
            barEntry = tempEntry;
        }
    }
    assertThat(childIt.hasNext()).isFalse();
    assertThat(attributeValues).containsExactlyInAnyOrder("foo1", "bar");

    assertThat(barEntry).isNotNull();
    // root: [foo], child: [foo1, bar1]
    barEntry.setValue("bar1");
    assertThat(child.attr(bar)).isEqualTo("bar1");

    // Set a new attribute to root.
    final AttributeKey<String> baz = AttributeKey.valueOf("baz");
    // root: [foo, baz], child: [foo1, bar1]
    root.setAttr(baz, "baz");

    childIt = child.attrs();
    attributeValues.clear();
    attributeValues.add((String) childIt.next().getValue());
    attributeValues.add((String) childIt.next().getValue());
    assertThat(attributeValues).containsExactlyInAnyOrder("foo1", "bar1");

    assertThat(childIt.next().getValue()).isEqualTo("baz");
    // childIt does not yield foo in root.
    assertThat(childIt.hasNext()).isFalse();

    // child own attrs()
    attributeValues.clear();
    final Iterator<Entry<AttributeKey<?>, Object>> childOwnIt = child.ownAttrs();
    for (int i = 0; i < 2; i++) {
        attributeValues.add((String) childOwnIt.next().getValue());
    }
    assertThat(childOwnIt.hasNext()).isFalse();
    assertThat(attributeValues).containsExactlyInAnyOrder("foo1", "bar1");
}