connpool.h

Go to the documentation of this file.

/*    Copyright 2009 10gen Inc.
 *
 *    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.
 */

#pragma once

#include <stack>
#include "dbclient.h"
#include "redef_macros.h"

#include "../util/background.h"

namespace mongo {

    class Shard;
    class DBConnectionPool;

    class PoolForHost {
    public:
        PoolForHost()
            : _created(0) {}

        PoolForHost( const PoolForHost& other ) {
            assert(other._pool.size() == 0);
            _created = other._created;
            assert( _created == 0 );
        }

        ~PoolForHost();

        int numAvailable() const { return (int)_pool.size(); }

        void createdOne( DBClientBase * base );
        long long numCreated() const { return _created; }

        ConnectionString::ConnectionType type() const { assert(_created); return _type; }

        DBClientBase * get( DBConnectionPool * pool , double socketTimeout );

        void done( DBConnectionPool * pool , DBClientBase * c );

        void flush();
        
        void getStaleConnections( vector<DBClientBase*>& stale );

        static void setMaxPerHost( unsigned max ) { _maxPerHost = max; }
        static unsigned getMaxPerHost() { return _maxPerHost; }
    private:

        struct StoredConnection {
            StoredConnection( DBClientBase * c );

            bool ok( time_t now );

            DBClientBase* conn;
            time_t when;
        };

        std::stack<StoredConnection> _pool;
        
        long long _created;
        ConnectionString::ConnectionType _type;

        static unsigned _maxPerHost;
    };

    class DBConnectionHook {
    public:
        virtual ~DBConnectionHook() {}
        virtual void onCreate( DBClientBase * conn ) {}
        virtual void onHandedOut( DBClientBase * conn ) {}
        virtual void onDestory( DBClientBase * conn ) {}
    };

    class DBConnectionPool : public PeriodicTask {
        
    public:

        DBConnectionPool();
        ~DBConnectionPool();

        void setName( const string& name ) { _name = name; }

        void onCreate( DBClientBase * conn );
        void onHandedOut( DBClientBase * conn );
        void onDestory( DBClientBase * conn );

        void flush();

        DBClientBase *get(const string& host, double socketTimeout = 0);
        DBClientBase *get(const ConnectionString& host, double socketTimeout = 0);

        void release(const string& host, DBClientBase *c);

        void addHook( DBConnectionHook * hook ); // we take ownership
        void appendInfo( BSONObjBuilder& b );

        struct serverNameCompare {
            bool operator()( const string& a , const string& b ) const;
        };

        virtual string taskName() const { return "DBConnectionPool-cleaner"; }
        virtual void taskDoWork();        

    private:
        DBConnectionPool( DBConnectionPool& p );
        
        DBClientBase* _get( const string& ident , double socketTimeout );

        DBClientBase* _finishCreate( const string& ident , double socketTimeout, DBClientBase* conn );
        
        struct PoolKey {
            PoolKey( string i , double t ) : ident( i ) , timeout( t ) {}
            string ident;
            double timeout;
        };

        struct poolKeyCompare {
            bool operator()( const PoolKey& a , const PoolKey& b ) const;
        };

        typedef map<PoolKey,PoolForHost,poolKeyCompare> PoolMap; // servername -> pool

        mongo::mutex _mutex;
        string _name;
        
        PoolMap _pools;

        // pointers owned by me, right now they leak on shutdown
        // _hooks itself also leaks because it creates a shutdown race condition
        list<DBConnectionHook*> * _hooks; 

    };

    extern DBConnectionPool pool;

    class AScopedConnection : boost::noncopyable {
    public:
        AScopedConnection() { _numConnections++; }
        virtual ~AScopedConnection() { _numConnections--; }
        
        virtual DBClientBase* get() = 0;
        virtual void done() = 0;
        virtual string getHost() const = 0;
        
        virtual bool ok() const = 0;

        static int getNumConnections() { return _numConnections; }

    private:
        static AtomicUInt _numConnections;
    };

    class ScopedDbConnection : public AScopedConnection {
    public:
        explicit ScopedDbConnection(const string& host, double socketTimeout = 0) : _host(host), _conn( pool.get(host, socketTimeout) ), _socketTimeout( socketTimeout ) {
            _setSocketTimeout();
        }

        ScopedDbConnection() : _host( "" ) , _conn(0), _socketTimeout( 0 ) {}

        /* @param conn - bind to an existing connection */
        ScopedDbConnection(const string& host, DBClientBase* conn, double socketTimeout = 0 ) : _host( host ) , _conn( conn ), _socketTimeout( socketTimeout ) {
            _setSocketTimeout();
        }

        explicit ScopedDbConnection(const ConnectionString& url, double socketTimeout = 0 ) : _host(url.toString()), _conn( pool.get(url, socketTimeout) ), _socketTimeout( socketTimeout ) {
            _setSocketTimeout();
        }

        explicit ScopedDbConnection(const Shard& shard, double socketTimeout = 0 );
        explicit ScopedDbConnection(const Shard* shard, double socketTimeout = 0 );

        ~ScopedDbConnection();

        DBClientBase* operator->() {
            uassert( 11004 ,  "connection was returned to the pool already" , _conn );
            return _conn;
        }

        DBClientBase& conn() {
            uassert( 11005 ,  "connection was returned to the pool already" , _conn );
            return *_conn;
        }

        DBClientBase* get() {
            uassert( 13102 ,  "connection was returned to the pool already" , _conn );
            return _conn;
        }

        bool ok() const { return _conn > 0; }

        string getHost() const { return _host; }

        void kill() {
            delete _conn;
            _conn = 0;
        }

        void done() {
            if ( ! _conn )
                return;

            /* we could do this, but instead of assume one is using autoreconnect mode on the connection
            if ( _conn->isFailed() )
                kill();
            else
            */
            pool.release(_host, _conn);
            _conn = 0;
        }

        ScopedDbConnection * steal();

    private:

        void _setSocketTimeout();

        const string _host;
        DBClientBase *_conn;
        const double _socketTimeout;

    };

} // namespace mongo

#include "undef_macros.h"

---