grpc-java issue 5692 example

Issue5692.java

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.Nullable;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import io.grpc.Attributes;
import io.grpc.EquivalentAddressGroup;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.NameResolver;
import io.grpc.NameResolverProvider;
import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.Status;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthGrpc;
import io.grpc.services.HealthStatusManager;

public class Issue5692 {

    public static class SampleDnsNameResolverProvider extends NameResolverProvider {

        static final String SCHEME = "sample-dns";

        public String getDefaultScheme() {
            return SCHEME;
        }

        @Nullable
        @Override
        public NameResolver newNameResolver(URI targetUri, NameResolver.Helper helper) {
            if (SCHEME.equals(targetUri.getScheme())) {
                String targetPath = Preconditions.checkNotNull(targetUri.getPath(), "targetPath");
                Preconditions.checkArgument(targetPath.startsWith("/"),
                        "the path component (%s) of the target (%s) must start with '/'", targetPath, targetUri);
                String name = targetPath.substring(1);
                return new SampleDnsNameResolver(name, helper);
            }
            return null;
        }

        @Override
        protected boolean isAvailable() {
            return true;
        }

        @Override
        protected int priority() {
            return 5;
        }
    }

    public static class SampleDnsNameResolver extends NameResolver {

        private final ExecutorService executor = Executors.newSingleThreadExecutor(new ThreadFactoryBuilder().setDaemon(true).setNameFormat("sample-dns-resolver-%d").build());
        private final String authority;
        private final String name;
        private final int port;
        private Listener listener;
        private AtomicBoolean resolving = new AtomicBoolean();

        SampleDnsNameResolver(String name, Helper helper) {
            final URI nameUri = URI.create("//" + Objects.requireNonNull(name, "name"));
            this.authority = Objects.requireNonNull(nameUri.getAuthority());
            this.name = nameUri.getHost();
            if (nameUri.getPort() == -1) {
                port = helper.getDefaultPort();
            } else {
                port = nameUri.getPort();
            }
        }

        public String getServiceAuthority() {
            return authority;
        }

        public synchronized void start(Listener listener) {
            Preconditions.checkState(listener != null);
            this.listener = Objects.requireNonNull(listener);
            executor.submit(() -> {
                resolving.set(true);
                try {
                    // NOTE: First resolution, return empty list - this won't recover
                    listener.onAddresses(Collections.emptyList(), Attributes.EMPTY);

                    // NOTE: This works
                    // listener.onError(Status.UNAVAILABLE.withDescription("no addresses available"));
                } finally {
                    resolving.set(false);
                }
            });
        }

        @Override
        public void refresh() {
            if (resolving.compareAndSet(false, true)) {
                executor.submit(() -> {
                    try {
                        final InetAddress[] addresses = InetAddress.getAllByName(name);
                        if (addresses == null || addresses.length == 0) {
                            listener.onError(Status.UNAVAILABLE.withDescription("No addresses found for name: " + name));
                            return;
                        }
                        final List<EquivalentAddressGroup> groups = new ArrayList<>(addresses.length);
                        for (InetAddress address : addresses) {
                            groups.add(new EquivalentAddressGroup(new InetSocketAddress(address, port)));
                        }
                        listener.onAddresses(groups, Attributes.EMPTY);
                    } catch (UnknownHostException e) {
                        listener.onError(Status.UNAVAILABLE.withCause(e));
                    } finally {
                        resolving.set(false);
                    }
                });
            }
        }

        public void shutdown() {
            executor.shutdown();
        }
    }

    public static void main(String[] args) throws Exception {
        final HealthStatusManager manager = new HealthStatusManager();
        manager.setStatus(HealthStatusManager.SERVICE_NAME_ALL_SERVICES, HealthCheckResponse.ServingStatus.SERVING);
        final Server server = ServerBuilder.forPort(8080)
            .addService(manager.getHealthService())
            .build();
        server.start();
        final ManagedChannel channel = ManagedChannelBuilder.forTarget("sample-dns:///localhost:8080")
            .usePlaintext()
            .build();
        boolean succeeded = false;
        final HealthGrpc.HealthBlockingStub stub = HealthGrpc.newBlockingStub(channel);
        for (int i = 0; i < 10; i++) {
            try {
                if (stub.check(HealthCheckRequest.getDefaultInstance()).getStatus() == HealthCheckResponse.ServingStatus.SERVING) {
                    succeeded = true;
                    break;
                }
            } catch (Exception e) {
                System.err.println("Failed to call health service: " + e.getMessage());
                Thread.sleep(1_000);
            }
        }
        channel.shutdown();
        server.shutdown();
        server.awaitTermination(1, TimeUnit.SECONDS);
        if (!succeeded) {
            throw new IllegalStateException("Didn't get a successful status");
        }
        System.out.println("Successfully called health service!");
    }
}