Javaで複数のRange検索を比較的高速に行うコード

• 
• Tweet

複数のCIDR形式で記されたネットワークに該当するか、複数の定義された時間帯に現時刻が該当するか、などなどRange検索をしたい事がそれなりにあります。

単純なRange検索であれば、from <= v && v <= toのような単純なコードで判定できますが、Rangeの定義が複数ある場合、愚直に実装するとO(N)になります。それはとても嫌なのでO(log N)で判定できるようにします。全部自前で書いてやろうと思ったのですが、ライブラリに任せることにしました。

Guavaのコードを見ると、追加するたびにRangeをMergeするという方法を取っていました。あるクエリにHitする定義はどれだみたいな事は出来ないのですが、今回は含まれるか(contains)を判定できれば良いので妥協します。

実際のコード その1

ライブラリはデータ構造としてguavaのRangeSet、CIDRの扱いとしてcommons-netのSubnetUtilsを使います。コードの記述を減らすためにLombok、JSONの処理にはgsonを使いました。

pom.xml

<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>at.orz</groupId>
    <artifactId>sample-range-checker</artifactId>
    <version>1.0.0-SNAPSHOT</version>

    <name>sample-range-checker</name>

    <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <maven.compiler.source>11</maven.compiler.source>
        <maven.compiler.target>11</maven.compiler.target>
    </properties>

    <dependencies>
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>27.1-jre</version>
        </dependency>
        <dependency>
            <groupId>commons-net</groupId>
            <artifactId>commons-net</artifactId>
            <version>3.6</version>
        </dependency>
        <dependency>
            <groupId>com.google.code.gson</groupId>
            <artifactId>gson</artifactId>
            <version>2.8.5</version>
        </dependency>
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
            <version>1.18.6</version>
            <scope>provided</scope>
        </dependency>
    </dependencies>
</project>

App.java

public class App {

  @Data
  public static class AWSIpRanges implements Serializable {
    private String syncToken;
    private String createDate;
    private List<IPv4Prefix> prefixes;
    private List<IPv6Prefix> ipv6Prefixes;
    @Data
    public static class IPv4Prefix implements Serializable {
      private String ipPrefix;
      private String region;
      private String service;
    }
    @Data
    public static class IPv6Prefix implements Serializable {
      private String ipv6Prefix;
      private String region;
      private String service;
    }
  }

  public static class ImmutableIntRangeChecker {
    @Delegate
    private final ImmutableRangeSet<Integer> rangeSet;
    public ImmutableIntRangeChecker(List<Range<Integer>> ranges) {
      this.rangeSet = ImmutableRangeSet.unionOf(ranges);
    }
    public boolean isIPAddrContains(String ipAddr) {
      int query = InetAddresses.coerceToInteger(InetAddresses.forString(ipAddr));
      return rangeSet.contains(query);
    }
  }

  public static void main(String[] args) throws IOException, InterruptedException {

    URI awsIpRangesURI = URI.create("https://ip-ranges.amazonaws.com/ip-ranges.json");

    // get aws-ip json file
    HttpClient client = HttpClient.newHttpClient();
    HttpRequest request = HttpRequest.newBuilder(awsIpRangesURI).build();
    String json = client.send(request, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8)).body();

    // parse json
    Gson gson = new GsonBuilder().setFieldNamingPolicy(FieldNamingPolicy.LOWER_CASE_WITH_UNDERSCORES).create();
    AWSIpRanges ipRanges =  gson.fromJson(json, AWSIpRanges.class);

    List<Range<Integer>> ranges = ipRanges.getPrefixes().stream()
        .map(AWSIpRanges.IPv4Prefix::getIpPrefix) // List<IP-String>にする
        .map(ip -> ip.contains("/") ? ip : ip + "/32") // SubnetUtilsでエラーにならないように
        .map(ip -> {
          // CIDR形式をパースし、ネットワークアドレスの範囲を求め、そのIPアドレスを32ビット変数に変換
          SubnetUtils subnet = new SubnetUtils(ip);
          subnet.setInclusiveHostCount(true);
          SubnetUtils.SubnetInfo info = subnet.getInfo();
          return Range.closed(info.asInteger(info.getLowAddress()), info.asInteger(info.getHighAddress()));
        })
        .collect(Collectors.toList());

    ImmutableIntRangeChecker checker = new ImmutableIntRangeChecker(ranges);
    System.out.println(checker.isIPAddrContains("35.172.155.127")); // -> true
    System.out.println(checker.isIPAddrContains("8.8.8.8"));        // -> false

  }
}

実際のコード その2

Intで扱わない場合です。RangeクラスはComparatorの実装を要求しますがInetAddrクラスは実装していないので、ラッパークラスを作成します。

またCheckerクラスを汎用的に実装し、FunctionクラスによってKeyの変換も定義します。今回はStringからComparableInetAddrに変換する関数を定義します。

App2.java

public static class ComparableInetAddr implements Serializable, Comparable<ComparableInetAddr> {
  @Delegate
  private final InetAddress inetAddress;
  private final int intIP;
  private ComparableInetAddr(InetAddress inetAddress) {
    this.inetAddress = inetAddress;
    this.intIP = InetAddresses.coerceToInteger(inetAddress);
  }
  public ComparableInetAddr(String ipAddr) {
    this(InetAddresses.forString(ipAddr));
  }
  @Override
  public int compareTo(ComparableInetAddr o) {
    return Integer.compare(this.intIP, o.intIP);
  }
}

public static class ImmutableRangeChecker<T, C extends Comparable> {
  @Delegate
  private final ImmutableRangeSet<C> rangeSet;
  private final Function<T, C> queryTransformer;
  public ImmutableRangeChecker(List<Range<C>> ranges, Function<T, C> queryTransformer) {
    this.rangeSet = ImmutableRangeSet.unionOf(ranges);
    this.queryTransformer = queryTransformer;
  }
  public boolean contains(T value) {
    C query = queryTransformer.apply(value);
    return rangeSet.contains(query);
  }
}

差分のコードだけ。ComparableIntAddrをRangeクラスに入れるようにした

List<Range<ComparableInetAddr>> ranges = ipRanges.getPrefixes().stream()
    .map(AWSIpRanges.IPv4Prefix::getIpPrefix)
    .map(ip -> ip.contains("/") ? ip : ip + "/32")
    .map(ip -> {
      SubnetUtils subnet = new SubnetUtils(ip);
      subnet.setInclusiveHostCount(true);
      SubnetUtils.SubnetInfo info = subnet.getInfo();
      // ここを変更した
      return Range.closed(new ComparableInetAddr(info.getLowAddress()), new ComparableInetAddr(info.getHighAddress()));
    })
    .collect(Collectors.toList());

ImmutableRangeChecker<String, ComparableInetAddr> checker =
    new ImmutableRangeChecker<>(
        ranges,
        ipString -> new ComparableInetAddr(InetAddresses.forString(ipString)));

System.out.println(checker.contains("35.172.155.127")); // false
System.out.println(checker.contains("8.8.8.8"));        // true

時間の範囲を判定する場合

例えば、1日のうちの10時から12時と14時から16時。とか。1週間のうちの月曜日の午前中と火曜日の午前中とか。前者の場合は0から86400の整数として扱えば良いし、後者は月曜始まりの0時を0とした整数として扱うとかすれば良いと思います。

もしくは上記のコードでComparableが実装されているDate系のクラスを使えば良いと思います。

その他

今回はImmutableRangeSetを使いましたが、他にもTreeRangeSetの実装があります。こっちはInstanceを生成後にRangeを追加することができます。個人的にはIntegerではなくPrimitiveのintで扱いたいです。

RangeがMergeされるのを確認したコード

@Test
public void rangeSample1() {
  ImmutableList<Range<Integer>> x = ImmutableList.of(Range.open(1, 100), Range.open(200, 300)).asList();
  TreeRangeSet<Integer> im = TreeRangeSet.create();
  im.addAll(x);
  im.add(Range.open(50, 250));
  System.out.println(im); // [(1..300)]
}
@Test
public void rangeSample2() {
  ImmutableList<Range<Integer>> x = ImmutableList.of(Range.closed(1, 100), Range.closed(200, 300)).asList();
  TreeRangeSet<Integer> im = TreeRangeSet.create();
  im.addAll(x);
  im.add(Range.open(50, 150)); // [[1..150), [200..300]]
  System.out.println(im);
}

参考

CIDRに限ってしまえばもっと高速にできるみたいです。

• http://dsas.blog.klab.org/archives/51293334.html
• https://stackoverflow.com/questions/45065994/create-guava-immutablerangeset-from-overlapping-ranges