001/*
002 * Copyright 2015-2018 the original author or authors
003 *
004 * This software is licensed under the Apache License, Version 2.0,
005 * the GNU Lesser General Public License version 2 or later ("LGPL")
006 * and the WTFPL.
007 * You may choose either license to govern your use of this software only
008 * upon the condition that you accept all of the terms of either
009 * the Apache License 2.0, the LGPL 2.1+ or the WTFPL.
010 */
011package org.minidns.source;
012
013import java.io.IOException;
014import java.net.InetAddress;
015import java.util.Locale;
016import java.util.concurrent.atomic.AtomicInteger;
017
018import org.minidns.AbstractDnsClient;
019import org.minidns.dnsmessage.DnsMessage;
020
021public class NetworkDataSourceWithAccounting extends NetworkDataSource {
022
023    private final AtomicInteger successfulQueries = new AtomicInteger();
024    private final AtomicInteger responseSize = new AtomicInteger();
025    private final AtomicInteger failedQueries = new AtomicInteger();
026
027    private final AtomicInteger successfulUdpQueries = new AtomicInteger();
028    private final AtomicInteger udpResponseSize = new AtomicInteger();
029    private final AtomicInteger failedUdpQueries = new AtomicInteger();
030
031    private final AtomicInteger successfulTcpQueries = new AtomicInteger();
032    private final AtomicInteger tcpResponseSize = new AtomicInteger();
033    private final AtomicInteger failedTcpQueries = new AtomicInteger();
034
035    @Override
036    public DnsMessage query(DnsMessage message, InetAddress address, int port) throws IOException {
037        DnsMessage response;
038        try {
039            response = super.query(message, address, port);
040        } catch (IOException e) {
041            failedQueries.incrementAndGet();
042            throw e;
043        }
044
045        successfulQueries.incrementAndGet();
046        responseSize.addAndGet(response.toArray().length);
047
048        return response;
049    }
050
051    @Override
052    protected DnsMessage queryUdp(DnsMessage message, InetAddress address, int port) throws IOException {
053        DnsMessage response;
054        try {
055            response = super.queryUdp(message, address, port);
056        } catch (IOException e) {
057            failedUdpQueries.incrementAndGet();
058            throw e;
059        }
060
061        successfulUdpQueries.incrementAndGet();
062        udpResponseSize.addAndGet(response.toArray().length);
063
064        return response;
065    }
066
067    @Override
068    protected DnsMessage queryTcp(DnsMessage message, InetAddress address, int port) throws IOException {
069        DnsMessage response;
070        try {
071            response = super.queryTcp(message, address, port);
072        } catch (IOException e) {
073            failedTcpQueries.incrementAndGet();
074            throw e;
075        }
076
077        successfulTcpQueries.incrementAndGet();
078        tcpResponseSize.addAndGet(response.toArray().length);
079
080        return response;
081    }
082
083    public Stats getStats() {
084        return new Stats(this);
085    }
086
087    public static NetworkDataSourceWithAccounting from(AbstractDnsClient client) {
088        DnsDataSource ds = client.getDataSource();
089        if (ds instanceof NetworkDataSourceWithAccounting) {
090            return (NetworkDataSourceWithAccounting) ds;
091        }
092        return null;
093    }
094
095    public static class Stats {
096        public final int successfulQueries;
097        public final int responseSize;
098        public final int averageResponseSize;
099        public final int failedQueries;
100
101        public final int successfulUdpQueries;
102        public final int udpResponseSize;
103        public final int averageUdpResponseSize;
104        public final int failedUdpQueries;
105
106        public final int successfulTcpQueries;
107        public final int tcpResponseSize;
108        public final int averageTcpResponseSize;
109        public final int failedTcpQueries;
110
111        private String stringCache;
112
113        private Stats(NetworkDataSourceWithAccounting ndswa) {
114            successfulQueries = ndswa.successfulQueries.get();
115            responseSize = ndswa.responseSize.get();
116            failedQueries = ndswa.failedQueries.get();
117
118            successfulUdpQueries = ndswa.successfulUdpQueries.get();
119            udpResponseSize = ndswa.udpResponseSize.get();
120            failedUdpQueries = ndswa.failedUdpQueries.get();
121
122            successfulTcpQueries = ndswa.successfulTcpQueries.get();
123            tcpResponseSize = ndswa.tcpResponseSize.get();
124            failedTcpQueries = ndswa.failedTcpQueries.get();
125
126            // Calculated stats section
127            averageResponseSize = successfulQueries > 0 ? responseSize / successfulQueries : 0;
128            averageUdpResponseSize = successfulUdpQueries > 0 ? udpResponseSize / successfulUdpQueries : 0;
129            averageTcpResponseSize = successfulTcpQueries > 0 ? tcpResponseSize / successfulTcpQueries : 0;
130        }
131
132        @Override
133        public String toString() {
134            if (stringCache != null)
135                return stringCache;
136
137            StringBuilder sb = new StringBuilder();
138
139            sb.append("Stats\t").append("# Successful").append('\t').append("# Failed").append('\t')
140                    .append("Resp. Size").append('\t').append("Avg. Resp. Size").append('\n');
141            sb.append("Total\t").append(toString(successfulQueries)).append('\t').append(toString(failedQueries))
142                    .append('\t').append(toString(responseSize)).append('\t').append(toString(averageResponseSize))
143                    .append('\n');
144            sb.append("UDP\t").append(toString(successfulUdpQueries)).append('\t').append(toString(failedUdpQueries))
145                    .append('\t').append(toString(udpResponseSize)).append('\t')
146                    .append(toString(averageUdpResponseSize)).append('\n');
147            sb.append("TCP\t").append(toString(successfulTcpQueries)).append('\t').append(toString(failedTcpQueries))
148                    .append('\t').append(toString(tcpResponseSize)).append('\t')
149                    .append(toString(averageTcpResponseSize)).append('\n');
150
151            stringCache = sb.toString();
152            return stringCache;
153        }
154
155        private static String toString(int i) {
156            return String.format(Locale.US, "%,09d", i);
157        }
158    }
159}