memcachepool.py

##
# Copyright (c) 2010 Mochi Media, Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##

from twisted.protocols.memcache import MemCacheProtocol

from txconnpool.pool import PooledClientFactory, Pool
from twisted.python import log


class _PooledMemCacheProtocol(MemCacheProtocol):
    """
    A MemCacheProtocol that will notify a connectionPool that it is ready
    to accept requests.
    """
    factory = None

    def connectionMade(self):
        """
        Notify our factory that we're ready to accept connections.
        """
        MemCacheProtocol.connectionMade(self)

        self.factory.connectionPool.clientFree(self)

        if self.factory.deferred is not None:
            self.factory.deferred.callback(self)
            self.factory.deferred = None


class _MemCacheClientFactory(PooledClientFactory):
    """
    L{PooledClientFactory} that uses L{_PooledMemCacheProtocol}
    """
    protocol = _PooledMemCacheProtocol



    def clientConnectionLost(self, connector, reason):
        log.msg("Client Connection Lost")
        PooledClientFactory.clientConnectionLost(
            self,
            connector,
            reason)


    def clientConnectionFailed(self, connector, reason):
        log.msg("Client Connection Failed")
        PooledClientFactory.clientConnectionFailed(
            self,
            connector,
            reason)


class MemCachePool(Pool):
    """
    A MemCache client which is backed by a pool of connections.

    Usage Example::
        from twisted.internet.address import IPv4Address
        from txconnpool.memcache import MemCachePool

        addr = IPv4Address('TCP', '127.0.0.1', 11211)
        mc_pool = MemCachePool(addr, maxClients=20)

        d = mc_pool.get('cached-data')

        def gotCachedData(data):
            flags, value = data
            if value:
                print 'Yay, we got a cache hit'
            else:
                print 'Boo, it was a cache miss'

        d.addCallback(gotCachedData)
    """
    clientFactory = _MemCacheClientFactory

    def get(self, *args, **kwargs):
        """
        See L{twisted.protocols.memcache.MemCacheProtocol.get}.
        """
        return self.performRequest('get', *args, **kwargs)

    def set(self, *args, **kwargs):
        """
        See L{twisted.protocols.memcache.MemCacheProtocol.set}
        """
        return self.performRequest('set', *args, **kwargs)

    def delete(self, *args, **kwargs):
        """
        See L{twisted.protocols.memcache.MemCacheProtocol.delete}
        """
        return self.performRequest('delete', *args, **kwargs)

    def add(self, *args, **kwargs):
        """
        See L{twisted.protocols.memcache.MemCacheProtocol.add}
        """
        return self.performRequest('add', *args, **kwargs)

txmemcachedns.py

# This is a simple implementation of a DNS server which looks up A records first in
# cache (with a restricted TTL), falling back to a Memcache lookup and finally
# a recursive DNS query.

# Original: 

# Python/Twisted/Redis backed DNS server - resolves from NAME to IP addrs
# fallback to google or any other DNS server to resolv domains not present on Redis
# to set a new domain on redis, just issue a SET domain.tld ip_addr
# run with twistd -ny txredns.tac
# gleicon 2011

from twisted.internet import reactor, task, protocol, defer
from twisted.names import dns, server, client, cache
from twisted.application import service, internet
from twisted.python.failure import Failure
from twisted.python import log
from twisted.internet.address import IPv4Address
from memcachepool import MemCachePool
import re


class memcache_resolver_backend(client.Resolver):
    def __init__(self, memcache, servers=None):
        self.memcache = memcache
        client.Resolver.__init__(self, servers=servers)
        self.ttl = 15
        self.prefix = "_dns_lookup_"

    def _handle_local_override(self, hostname):
        return [(dns.RRHeader(hostname, dns.A, dns.IN, self.ttl, dns.Record_A("127.0.0.1", self.ttl)),), (), ()]

    @defer.inlineCallbacks
    def _get_ip_addr(self, hostname, timeout):
        (a, ip) = yield self.memcache.get(self.prefix + hostname)
        log.msg('looked up %s in memcache: %s'% (hostname, ip))
        if ip:
            defer.returnValue([(dns.RRHeader(hostname, dns.A, dns.IN, self.ttl, dns.Record_A(ip, self.ttl)),), (), ()])
        else:
            i = yield self._lookup(hostname, dns.IN, dns.A, timeout)

            # Force a lower TTL than the default.
            if ( len(i) > 0 and len(i[0]) > 0 and i[0][0].__dict__.has_key('ttl') ):
                i[0][0].ttl = self.ttl

            defer.returnValue(i)

    def lookupAddress(self, hostname, timeout = None):
        if (re.match( r'.+\.dev', hostname )):
            return self._handle_local_override(hostname)

        return self._get_ip_addr(hostname, timeout)
       
def create_application():
    memcache_address = IPv4Address('TCP', '127.0.0.1', 11211)
    memcache_pool = MemCachePool(memcache_address, maxClients=20)

    resolver = memcache_resolver_backend(memcache_pool, servers=[('8.8.8.8', 53), ('8.8.4.4', '53')])

    application = service.Application("txmemcachedns")
    srv_collection = service.IServiceCollection(application)

    dnsFactory = server.DNSServerFactory(caches=[cache.CacheResolver()], clients=[resolver])

    internet.TCPServer(53, dnsFactory).setServiceParent(srv_collection)
    internet.UDPServer(53, dns.DNSDatagramProtocol(dnsFactory)).setServiceParent(srv_collection)
    return application

# .tac app
application = create_application()
	

This has been tested for A records, but not MX, TXT, or anything else. So it may hose your mailserver. Use with caution and test first. :D