com.clickzetta.platform.test.MockServer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation
The java SDK for clickzetta's Lakehouse
There is a newer version: 2.0.0
package com.clickzetta.platform.test;

import com.clickzetta.platform.connection.ChannelManager;
import com.google.common.annotations.VisibleForTesting;
import io.grpc.*;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@VisibleForTesting
public class MockServer {

  private static final Logger LOG = LoggerFactory.getLogger(MockServer.class);

  private int port;
  private Server server;

  // grpc service instance.
  private final List serviceBases;

  public MockServer(int port, List serviceBases) {
    this.port = port;
    this.serviceBases = serviceBases;
  }

  public List getServiceBases() {
    return serviceBases;
  }

  public void start() throws IOException {
    synchronized (this) {
      if (server != null) {
        return;
      }
      server = ServerBuilder.forPort(port)
          .addServices(serviceBases.stream().map(BindableService::bindService).collect(Collectors.toList()))
          .intercept(new ServerInterceptor() {
            @Override
            public 
            ServerCall.Listener interceptCall(ServerCall serverCall, Metadata metadata,
                                                    ServerCallHandler serverCallHandler) {
              Metadata.Key streamId = Metadata.Key.of(ChannelManager.RPC_CLIENT_STREAM_ID, Metadata.ASCII_STRING_MARSHALLER);
              if (metadata.containsKey(streamId)) {
                LOG.info("grpc server intercept call receive streamId: {}", metadata.get(streamId));
              }
              return serverCallHandler.startCall(serverCall, metadata);
            }
          })
          .maxInboundMessageSize(20 * 1024 * 1024)
          .maxInboundMetadataSize(4 * 1024 * 1024)
          .build()
          .start();
    }
    Runtime.getRuntime().addShutdownHook(new Thread(() -> {
      try {
        MockServer.this.stop();
      } catch (IOException e) {
        e.printStackTrace(System.err);
      }
      LOG.info("Mock Server {} is shutting down.", StringUtils.join(ClassUtil.getListClassName(serviceBases), ","));
    }));
    LOG.info("Mock Server {} started, listening on {}.", StringUtils.join(ClassUtil.getListClassName(serviceBases), ","), port);
    if (server != null) {
      try {
        server.awaitTermination();
      } catch (InterruptedException ite) {
        throw new IOException(ite);
      }
    }
  }

  public void stop() throws IOException {
    Server toCloseServer = null;
    synchronized (this) {
      if (server == null) {
        return;
      }
      toCloseServer = server;
      server = null;
    }
    if (toCloseServer != null) {
      try {
        toCloseServer.shutdown().awaitTermination(1 * 1000, TimeUnit.MILLISECONDS);
        if (toCloseServer != null && !toCloseServer.isShutdown()) {
          toCloseServer.shutdownNow();
        }
      } catch (InterruptedException ite) {
        throw new IOException(ite);
      }
    }
  }
}