Session Distribution

/**
 * Abstract base class for session distributors
 */
abstract class Distributor {
    /** Protected constructor with required dependencies */
    protected Distributor(Tracer tracer, HttpClient.Factory httpClientFactory, Secret registrationSecret);
    
    /** Create a new WebDriver session */
    abstract Either<SessionNotCreatedException, CreateSessionResponse> newSession(SessionRequest request);
    
    /** Add a node to the grid */
    abstract Distributor add(Node node);
    
    /** Drain sessions from a node (stop accepting new sessions) */
    abstract boolean drain(NodeId nodeId);
    
    /** Remove a node from the grid */
    abstract void remove(NodeId nodeId);
    
    /** Get current distributor status and statistics */
    abstract DistributorStatus getStatus();
}

/**
 * Local distributor implementation for single-process or hub deployments
 */
@ManagedService
class LocalDistributor extends Distributor {
    /** Factory method to create distributor from configuration */
    static Distributor create(Config config);
    
    // Implement all abstract methods from Distributor
    Either<SessionNotCreatedException, CreateSessionResponse> newSession(SessionRequest request);
    Distributor add(Node node);
    boolean drain(NodeId nodeId);
    void remove(NodeId nodeId);
    DistributorStatus getStatus();
    
    // JMX management attributes
    @ManagedAttribute
    int getNodeCount();
    
    @ManagedAttribute
    int getSessionCount();
    
    @ManagedAttribute
    long getSessionsCreated();
    
    @ManagedAttribute
    Map<String, Object> getSlotStereotypes();
}

Config config = new MemoizedConfig(new MapConfig(Map.of(
    "distributor", Map.of(
        "reject-unsupported-caps", true,
        "healthcheck-interval", "120s"
    )
)));

Distributor distributor = LocalDistributor.create(config);

// Add nodes to the distributor
Node chromeNode = createChromeNode();
Node firefoxNode = createFirefoxNode();

Distributor updatedDistributor = distributor
    .add(chromeNode)
    .add(firefoxNode);

// Create a new session
SessionRequest request = new SessionRequest(
    new RequestId(UUID.randomUUID()),
    Instant.now().plus(Duration.ofMinutes(5)), // timeout
    Set.of(W3C), // dialects
    new ImmutableCapabilities("browserName", "chrome")
);

Either<SessionNotCreatedException, CreateSessionResponse> result = 
    updatedDistributor.newSession(request);

if (result.isRight()) {
    CreateSessionResponse response = result.right();
    System.out.println("Session created: " + response.getSession().getId());
} else {
    SessionNotCreatedException error = result.left();
    System.err.println("Session creation failed: " + error.getMessage());
}

/**
 * Strategy interface for selecting appropriate slots for session requests
 */
@FunctionalInterface
interface SlotSelector {
    /**
     * Select slots that can handle the requested capabilities
     * @param capabilities The desired session capabilities
     * @param nodes Available nodes with their current status
     * @param slotMatcher Function to test capability matching
     * @return Set of slot IDs that can fulfill the request
     */
    Set<SlotId> selectSlot(Capabilities capabilities, Set<NodeStatus> nodes, SlotMatcher slotMatcher);
}

/**
 * Default slot selection implementation using round-robin with load balancing
 */
class DefaultSlotSelector implements SlotSelector {
    Set<SlotId> selectSlot(Capabilities capabilities, Set<NodeStatus> nodes, SlotMatcher slotMatcher);
}

// Custom selector that prefers nodes with lower load
SlotSelector loadBasedSelector = (capabilities, nodes, slotMatcher) -> {
    return nodes.stream()
        .filter(node -> node.hasCapacity(capabilities, slotMatcher))
        .sorted(Comparator.comparing(NodeStatus::getLoad))
        .flatMap(node -> node.getSlots().stream())
        .filter(slot -> slot.isSupporting(capabilities, slotMatcher))
        .map(Slot::getId)
        .limit(1)
        .collect(Collectors.toSet());
};

/**
 * Remote distributor client for distributed grid deployments
 */
class RemoteDistributor extends Distributor {
    RemoteDistributor(Tracer tracer, HttpClient.Factory httpClientFactory, 
                     URI distributorUri, Secret registrationSecret);
    
    // Implement all abstract methods via HTTP calls
    Either<SessionNotCreatedException, CreateSessionResponse> newSession(SessionRequest request);
    Distributor add(Node node);
    boolean drain(NodeId nodeId);
    void remove(NodeId nodeId);
    DistributorStatus getStatus();
}

/**
 * Configuration options for distributor behavior
 */
class DistributorOptions {
    static final String DISTRIBUTOR_SECTION = "distributor";
    
    /** Get implementation class for the distributor */
    String getDistributorImplementation(Config config);
    
    /** Get health check interval for nodes */
    Duration getHealthCheckInterval(Config config);
    
    /** Whether to reject requests with unsupported capabilities */
    boolean shouldRejectUnsupportedCaps(Config config);
    
    /** Get slot matcher implementation */
    SlotMatcher getSlotMatcher(Config config);
    
    /** Get slot selector implementation */
    SlotSelector getSlotSelector(Config config);
}

/**
 * Command-line flags for distributor configuration
 */
class DistributorFlags {
    @Parameter(names = {"--reject-unsupported-caps"}, 
               description = "Reject session requests with unsupported capabilities")
    Boolean rejectUnsupportedCaps;
    
    @Parameter(names = {"--healthcheck-interval"}, 
               description = "How often to check node health")
    int healthCheckInterval = 120;
    
    @Parameter(names = {"--slot-matcher"}, 
               description = "Full class name of SlotMatcher implementation")
    String slotMatcher = "org.openqa.selenium.grid.data.DefaultSlotMatcher";
    
    @Parameter(names = {"--slot-selector"}, 
               description = "Full class name of SlotSelector implementation")  
    String slotSelector = "org.openqa.selenium.grid.distributor.selector.DefaultSlotSelector";
}

/**
 * Request to create a new WebDriver session
 */
class SessionRequest {
    SessionRequest(RequestId requestId, Instant enqueued, Set<Dialect> downstreamDialects, Capabilities desiredCapabilities);
    RequestId getRequestId();
    Instant getEnqueued();
    Set<Dialect> getDownstreamDialects(); 
    Capabilities getDesiredCapabilities();
}

/**
 * Response containing created session information
 */
class CreateSessionResponse {
    CreateSessionResponse(Session session, byte[] downstreamEncodedResponse);
    Session getSession();
    byte[] getDownstreamEncodedResponse();
}

/**
 * Current status of the distributor
 */
class DistributorStatus {
    DistributorStatus(Set<NodeStatus> nodes);
    Set<NodeStatus> getNodes();
    boolean hasCapacity();
    boolean hasCapacity(Capabilities capabilities);
}

// Session creation can fail for several reasons
Either<SessionNotCreatedException, CreateSessionResponse> result = distributor.newSession(request);

if (result.isLeft()) {
    SessionNotCreatedException error = result.left();
    
    // Common reasons for failure:
    // - No nodes available with matching capabilities
    // - All matching nodes at capacity
    // - Request timeout exceeded
    // - Node health check failure
    
    System.err.println("Session creation failed: " + error.getMessage());
}