feat: add workers support

Author: duanhongyi

smartdns/sdns.py

@@ -3,9 +3,6 @@
 # version
 __version__ = '1.0.1.1'
 
-import socket
-import time
-import signal
 import sys
 import os
 import yaml
@@ -18,6 +15,7 @@
 from twisted.application import service, internet
 from twisted.names import dns, server, client, cache, common, resolve
 from twisted.internet import reactor
+from multiprocessing import cpu_count, Process
 from . import dnsserver, ippool, monitor
 
 
@@ -32,74 +30,51 @@ def loadconfig(path):
         return yaml.load(f, Loader=yaml.FullLoader)
 
 
-def get_local_ip():
-    import sys
-    import socket
-    import fcntl
-    import array
-    import struct
-    is_64bits = sys.maxsize > 2**32
-    struct_size = 40 if is_64bits else 32
-    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    max_possible = 8  # initial value
-    while True:
-        bytes = max_possible * struct_size
-        names = array.array('B', b'\0' * bytes)
-        outbytes = struct.unpack('iL', fcntl.ioctl(
-            s.fileno(),
-            0x8912,  # SIOCGIFCONF
-            struct.pack('iL', bytes, names.buffer_info()[0])
-        ))[0]
-        if outbytes == bytes:
-            max_possible *= 2
-        else:
-            break
-    namestr = names.tostring()
-    return [(namestr[i:i+16].split(b'\0', 1)[0],
-             socket.inet_ntoa(namestr[i+20:i+24]))
-            for i in range(0, outbytes, struct_size)]
-
-
 def prepare_run(run_env):
-        # load main config
+    # load main config
     logger.info('start to load %s ......' % run_env['conf'])
     conf = loadconfig(os.path.join(run_env['conf'], 'sdns.yaml'))
-
     logger.info('start to load A,SOA,NS record ......')
-    Amapping = loadconfig(os.path.join(run_env['conf'], 'a.yaml'))
-    NSmapping = loadconfig(os.path.join(run_env['conf'], 'ns.yaml'))
-    SOAmapping = loadconfig(os.path.join(run_env['conf'], 'soa.yaml'))
+    a_mapping = loadconfig(os.path.join(run_env['conf'], 'a.yaml'))
+    ns_mapping = loadconfig(os.path.join(run_env['conf'], 'ns.yaml'))
+    soa_mapping = loadconfig(os.path.join(run_env['conf'], 'soa.yaml'))
 
     # start monitor
     monitor_config = loadconfig(os.path.join(run_env['conf'], 'monitor.yaml'))
-    monitor_mapping = monitor.MonitorMapping(monitor_config, Amapping)
+    monitor_mapping = monitor.MonitorMapping(monitor_config, a_mapping)
     # load dns record config file
     logger.info('start to init IP pool ......')
-    Finder = ippool.IPPool(
+    finder = ippool.IPPool(
         os.path.join(run_env['conf'], 'ip.csv'),
         os.path.join(run_env['conf'], 'a.yaml'),
         monitor_mapping)
 
-    run_env['finder'] = Finder
-
-    listen_tcp_port = conf['listen']['tcp']
-    listen_udp_port = conf['listen']['udp']
+    run_env['finder'] = finder
 
     # set up a resolver that uses the mapping or a secondary nameserver
     dnsforward = []
-    for i in conf['dnsforward_ip']:
-        dnsforward.append((i, conf['dnsforward_port']))
+    for i in conf['dnsforward']:
+        dnsforward_ip, dnsforward_port = i.split(":")
+        dnsforward.append((dnsforward_ip, int(dnsforward_port)))
 
-    for ifc, ip in get_local_ip():
-        # create the protocols
-        SmartResolver = dnsserver.MapResolver(
-            Finder, Amapping, NSmapping, SOAmapping, servers=dnsforward)
+    # create the protocols
+    for listen_tcp in conf['listen']['tcp']:
+        listen_tcp_ip, listen_tcp_port = listen_tcp.split(":")
         f = dnsserver.SmartDNSFactory(
-            caches=[cache.CacheResolver()], clients=[SmartResolver])
-        p = dns.DNSDatagramProtocol(f)
-        f.noisy = p.noisy = False
-        run_env['tcp'].append([listen_tcp_port, f, ip])
-        run_env['udp'].append([listen_udp_port, p, ip])
+            caches=[cache.CacheResolver()], clients=[
+                dnsserver.MapResolver(
+                    finder, a_mapping, ns_mapping, soa_mapping, servers=dnsforward)])
+        f.noisy = False
+        run_env['tcp'].append([int(listen_tcp_port), f, listen_tcp_ip])
+    for listen_udp in conf['listen']['tcp']:
+        listen_udp_ip, listen_udp_port = listen_udp.split(":")
+        p = dns.DNSDatagramProtocol(dnsserver.SmartDNSFactory(
+            caches=[cache.CacheResolver()], clients=[
+                dnsserver.MapResolver(
+                    finder, a_mapping, ns_mapping, soa_mapping, servers=dnsforward)]))
+        p.noisy = False
+        run_env['udp'].append([int(listen_udp_port), p, listen_udp_ip])
+    return conf
 
 
 def main():
@@ -109,10 +84,16 @@ def main():
         run_env['conf'] = sys.argv[1]
     else:
         run_env['conf'] = '/etc/smartdns'
-        
-    prepare_run(run_env)
+
+    conf = prepare_run(run_env)
     for e in run_env['tcp']:
         reactor.listenTCP(e[0], e[1], interface=e[2])
     for e in run_env['udp']:
         reactor.listenUDP(e[0], e[1], interface=e[2])
-    reactor.run()
+    workers = []
+    for _ in range(conf.get("workers", cpu_count() * 2)):
+        process = Process(target=reactor.run)
+        process.daemon = True
+        process.start()
+        workers.append(process)
+    [worker.join() for worker in workers]

templates/sdns.yaml

@@ -1,13 +1,17 @@
 #Smart DNS main config file
 
+#进程数量
+workers: 4
+
 #监听端口
 listen:
-  tcp: 53
-  udp: 53
+  tcp:
+    - 0.0.0.0:53
+  udp:
+    - 0.0.0.0:53
 
 #dns forward
 #需要支持多个
-dnsforward_ip:
-  - 119.29.29.29
-  - 180.76.76.76
-dnsforward_port: 53
+dnsforward:
+  - 119.29.29.29:53
+  - 180.76.76.76:53