• Home
• org.seleniumhq.selenium
• selenium-session-map-redis
• 4.25.0
• source code
• RedisBackedSessionMap.java

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

org.openqa.selenium.grid.sessionmap.redis.RedisBackedSessionMap Maven / Gradle / Ivy

// Licensed to the Software Freedom Conservancy (SFC) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The SFC licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.openqa.selenium.grid.sessionmap.redis;

import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES_EVENT;
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID;
import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
import static org.openqa.selenium.remote.tracing.Tags.EXCEPTION;

import com.google.common.collect.ImmutableMap;
import io.lettuce.core.KeyValue;
import java.net.URI;
import java.net.URISyntaxException;
import java.time.Instant;
import java.util.List;
import java.util.logging.Logger;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.ImmutableCapabilities;
import org.openqa.selenium.NoSuchSessionException;
import org.openqa.selenium.events.EventBus;
import org.openqa.selenium.grid.config.Config;
import org.openqa.selenium.grid.data.NodeRemovedEvent;
import org.openqa.selenium.grid.data.NodeRestartedEvent;
import org.openqa.selenium.grid.data.Session;
import org.openqa.selenium.grid.data.SessionClosedEvent;
import org.openqa.selenium.grid.log.LoggingOptions;
import org.openqa.selenium.grid.server.EventBusOptions;
import org.openqa.selenium.grid.sessionmap.SessionMap;
import org.openqa.selenium.grid.sessionmap.config.SessionMapOptions;
import org.openqa.selenium.internal.Require;
import org.openqa.selenium.json.Json;
import org.openqa.selenium.redis.GridRedisClient;
import org.openqa.selenium.remote.SessionId;
import org.openqa.selenium.remote.tracing.AttributeKey;
import org.openqa.selenium.remote.tracing.AttributeMap;
import org.openqa.selenium.remote.tracing.Span;
import org.openqa.selenium.remote.tracing.Status;
import org.openqa.selenium.remote.tracing.Tracer;

public class RedisBackedSessionMap extends SessionMap {

  private static final Logger LOG = Logger.getLogger(RedisBackedSessionMap.class.getName());
  private static final Json JSON = new Json();
  private static final String REDIS_URI_KEY = "session.uri_key";
  private static final String REDIS_URI_VALUE = "session.uri_value";
  private static final String REDIS_CAPABILITIES_KEY = "session.capabilities_key";
  private static final String REDIS_CAPABILITIES_VALUE = "session.capabilities_value";
  private static final String REDIS_START_KEY = "session.start";
  private static final String REDIS_START_VALUE = "session.start_value";
  private static final String DATABASE_SYSTEM = AttributeKey.DATABASE_SYSTEM.getKey();
  private static final String DATABASE_OPERATION = AttributeKey.DATABASE_OPERATION.getKey();
  private final GridRedisClient connection;
  private final EventBus bus;
  private final URI serverUri;

  public RedisBackedSessionMap(Tracer tracer, URI serverUri, EventBus bus) {
    super(tracer);

    Require.nonNull("Redis Server Uri", serverUri);
    this.bus = Require.nonNull("Event bus", bus);
    this.connection = new GridRedisClient(serverUri);
    this.serverUri = serverUri;
    this.bus.addListener(SessionClosedEvent.listener(this::remove));

    this.bus.addListener(
        NodeRemovedEvent.listener(
            nodeStatus ->
                nodeStatus.getSlots().stream()
                    .filter(slot -> slot.getSession() != null)
                    .map(slot -> slot.getSession().getId())
                    .forEach(this::remove)));

    bus.addListener(
        NodeRestartedEvent.listener(nodeStatus -> this.removeByUri(nodeStatus.getExternalUri())));
  }

  public static SessionMap create(Config config) {
    Tracer tracer = new LoggingOptions(config).getTracer();
    EventBus bus = new EventBusOptions(config).getEventBus();
    URI sessionMapUri = new SessionMapOptions(config).getSessionMapUri();

    return new RedisBackedSessionMap(tracer, sessionMapUri, bus);
  }

  @Override
  public boolean add(Session session) {
    Require.nonNull("Session to add", session);

    try (Span span =
        tracer
            .getCurrentContext()
            .createSpan("MSET sessionUriKey  capabilitiesKey  ")) {
      AttributeMap attributeMap = tracer.createAttributeMap();
      SESSION_ID.accept(span, session.getId());
      SESSION_ID_EVENT.accept(attributeMap, session.getId());
      CAPABILITIES.accept(span, session.getCapabilities());
      CAPABILITIES_EVENT.accept(attributeMap, session.getCapabilities());
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);

      String uriKey = uriKey(session.getId());
      String uriValue = session.getUri().toString();
      String stereotypeKey = stereotypeKey(session.getId());
      String stereotypeJson = JSON.toJson(session.getStereotype());
      String capabilitiesKey = capabilitiesKey(session.getId());
      String capabilitiesJson = JSON.toJson(session.getCapabilities());
      String startKey = startKey(session.getId());
      String startValue = JSON.toJson(session.getStartTime());

      span.setAttribute(REDIS_URI_KEY, uriKey);
      span.setAttribute(REDIS_URI_VALUE, uriValue);
      span.setAttribute(REDIS_CAPABILITIES_KEY, capabilitiesKey);
      span.setAttribute(REDIS_CAPABILITIES_VALUE, capabilitiesJson);
      span.setAttribute(DATABASE_OPERATION, "MSET");
      attributeMap.put(REDIS_URI_KEY, uriKey);
      attributeMap.put(REDIS_URI_VALUE, uriValue);
      attributeMap.put(REDIS_CAPABILITIES_KEY, capabilitiesKey);
      attributeMap.put(REDIS_CAPABILITIES_VALUE, capabilitiesJson);
      attributeMap.put(REDIS_START_KEY, startKey);
      attributeMap.put(REDIS_START_VALUE, startValue);
      attributeMap.put(DATABASE_OPERATION, "MSET");

      span.addEvent("Inserted into the database", attributeMap);
      connection.mset(
          ImmutableMap.of(
              uriKey, uriValue,
              stereotypeKey, stereotypeJson,
              capabilitiesKey, capabilitiesJson,
              startKey, startValue));

      return true;
    }
  }

  @Override
  public Session get(SessionId id) throws NoSuchSessionException {
    Require.nonNull("Session ID", id);

    try (Span span = tracer.getCurrentContext().createSpan("GET capabilitiesKey")) {
      AttributeMap attributeMap = tracer.createAttributeMap();
      SESSION_ID.accept(span, id);
      SESSION_ID_EVENT.accept(attributeMap, id);
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);
      span.setAttribute(DATABASE_OPERATION, "GET");
      attributeMap.put(DATABASE_OPERATION, "GET");

      URI uri = getUri(id);

      attributeMap.put(REDIS_URI_KEY, uriKey(id));
      attributeMap.put(AttributeKey.SESSION_URI.getKey(), uri.toString());

      String capabilitiesKey = capabilitiesKey(id);
      String rawCapabilities = connection.get(capabilitiesKey);

      String stereotypeKey = stereotypeKey(id);
      String rawStereotype = connection.get(stereotypeKey);

      span.setAttribute(REDIS_CAPABILITIES_KEY, capabilitiesKey);
      attributeMap.put(REDIS_CAPABILITIES_KEY, capabilitiesKey);

      if (rawCapabilities != null) {
        span.setAttribute(REDIS_CAPABILITIES_VALUE, rawCapabilities);
      }

      Capabilities caps =
          rawCapabilities == null
              ? new ImmutableCapabilities()
              : JSON.toType(rawCapabilities, Capabilities.class);

      Capabilities stereotype =
          rawStereotype == null
              ? new ImmutableCapabilities()
              : JSON.toType(rawStereotype, Capabilities.class);

      String rawStart = connection.get(startKey(id));
      Instant start = JSON.toType(rawStart, Instant.class);

      CAPABILITIES.accept(span, caps);
      CAPABILITIES_EVENT.accept(attributeMap, caps);

      span.addEvent("Retrieved session from the database", attributeMap);
      return new Session(id, uri, stereotype, caps, start);
    }
  }

  @Override
  public URI getUri(SessionId id) throws NoSuchSessionException {
    Require.nonNull("Session ID", id);

    try (Span span = tracer.getCurrentContext().createSpan("GET sessionURI")) {
      AttributeMap attributeMap = tracer.createAttributeMap();
      SESSION_ID.accept(span, id);
      SESSION_ID_EVENT.accept(attributeMap, id);
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);
      span.setAttribute(DATABASE_OPERATION, "GET");
      attributeMap.put(DATABASE_OPERATION, "GET");

      String uriKey = uriKey(id);
      List> rawValues = connection.mget(uriKey);

      String rawUri = rawValues.get(0).getValueOrElse(null);

      span.setAttribute(REDIS_URI_KEY, uriKey);
      attributeMap.put(REDIS_URI_KEY, uriKey);

      if (rawUri == null) {
        NoSuchSessionException exception = new NoSuchSessionException("Unable to find session.");
        span.setAttribute("error", true);
        span.setStatus(Status.NOT_FOUND);
        EXCEPTION.accept(attributeMap, exception);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            "Session URI does not exist in the database :" + exception.getMessage());
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

        throw exception;
      }

      span.setAttribute(REDIS_URI_VALUE, rawUri);
      attributeMap.put(REDIS_URI_KEY, uriKey);
      attributeMap.put(REDIS_URI_VALUE, rawUri);

      try {
        return new URI(rawUri);
      } catch (URISyntaxException e) {
        span.setAttribute("error", true);
        span.setStatus(Status.INTERNAL);
        EXCEPTION.accept(attributeMap, e);
        attributeMap.put(AttributeKey.SESSION_URI.getKey(), rawUri);
        attributeMap.put(
            AttributeKey.EXCEPTION_MESSAGE.getKey(),
            "Unable to convert session id to uri: " + e.getMessage());
        span.addEvent(AttributeKey.EXCEPTION_EVENT.getKey(), attributeMap);

        throw new NoSuchSessionException(
            String.format("Unable to convert session id (%s) to uri: %s", id, rawUri), e);
      }
    }
  }

  @Override
  public void remove(SessionId id) {
    Require.nonNull("Session ID", id);

    try (Span span = tracer.getCurrentContext().createSpan("DEL sessionUriKey capabilitiesKey")) {
      AttributeMap attributeMap = tracer.createAttributeMap();
      SESSION_ID.accept(span, id);
      SESSION_ID_EVENT.accept(attributeMap, id);
      setCommonSpanAttributes(span);
      setCommonEventAttributes(attributeMap);
      span.setAttribute(DATABASE_OPERATION, "DEL");
      attributeMap.put(DATABASE_OPERATION, "DEL");

      String uriKey = uriKey(id);
      String capabilitiesKey = capabilitiesKey(id);
      String stereotypeKey = stereotypeKey(id);
      String startKey = startKey(id);
      span.setAttribute(REDIS_URI_KEY, uriKey);
      span.setAttribute(REDIS_CAPABILITIES_KEY, capabilitiesKey);
      span.setAttribute(REDIS_START_KEY, startKey);
      attributeMap.put(REDIS_URI_KEY, uriKey);
      attributeMap.put(REDIS_CAPABILITIES_KEY, capabilitiesKey);
      attributeMap.put(REDIS_START_KEY, startKey);

      span.addEvent("Deleted session from the database", attributeMap);

      connection.del(uriKey, capabilitiesKey, stereotypeKey, startKey);
    }
  }

  public void removeByUri(URI uri) {
    List uriKeys = connection.getKeysByPattern("session:*:uri");

    if (uriKeys.isEmpty()) {
      return;
    }

    String[] keys = new String[uriKeys.size()];
    keys = uriKeys.toArray(keys);

    List> keyValues = connection.mget(keys);
    keyValues.stream()
        .filter(entry -> entry.getValue().equals(uri.toString()))
        .map(KeyValue::getKey)
        .map(
            key -> {
              String[] sessionKey = key.split(":");
              return new SessionId(sessionKey[1]);
            })
        .forEach(this::remove);
  }

  @Override
  public boolean isReady() {
    return connection.isOpen();
  }

  private String uriKey(SessionId id) {
    Require.nonNull("Session ID", id);

    return "session:" + id + ":uri";
  }

  private String capabilitiesKey(SessionId id) {
    Require.nonNull("Session ID", id);

    return "session:" + id + ":capabilities";
  }

  private String startKey(SessionId id) {
    Require.nonNull("Session ID", id);

    return "session:" + id + ":start";
  }

  private String stereotypeKey(SessionId id) {
    Require.nonNull("Session ID", id);

    return "session:" + id + ":stereotype";
  }

  private void setCommonSpanAttributes(Span span) {
    span.setAttribute("span.kind", Span.Kind.CLIENT.toString());
    span.setAttribute(DATABASE_SYSTEM, "redis");
  }

  private void setCommonEventAttributes(AttributeMap map) {
    map.put(DATABASE_SYSTEM, "redis");
    if (serverUri != null) {
      map.put(AttributeKey.DATABASE_CONNECTION_STRING.getKey(), serverUri.toString());
    }
  }

  public GridRedisClient getRedisClient() {
    return connection;
  }
}