MemcachedClient.php

Go to the documentation of this file.

<?php
// phpcs:ignoreFile -- It's an external lib and it isn't. Let's not bother.
use Psr\Log\LoggerInterface;
use Psr\Log\NullLogger;
use Wikimedia\IPUtils;
 
// {{{ class MemcachedClient
class MemcachedClient {
    // {{{ properties
    // {{{ public
 
    // {{{ constants
    // {{{ flags
 
    const SERIALIZED = 1;
 
    const COMPRESSED = 2;
 
    const INTVAL = 4;
 
    // }}}
 
    const COMPRESSION_SAVINGS = 0.20;
 
    // }}}
 
    public $stats;
 
    // }}}
    // {{{ private
 
    public $_cache_sock;
 
    public $_debug;
 
    public $_host_dead;
 
    public $_have_zlib;
 
    public $_compress_enable;
 
    public $_compress_threshold;
 
    public $_persistent;
 
    public $_single_sock;
 
    public $_servers;
 
    public $_buckets;
 
    public $_bucketcount;
 
    public $_active;
 
    public $_timeout_seconds;
 
    public $_timeout_microseconds;
 
    public $_connect_timeout;
 
    public $_connect_attempts;
 
    private $_logger;
 
    // }}}
    // }}}
    // {{{ methods
    // {{{ public functions
    // {{{ memcached()
 
    public function __construct( $args ) {
        $this->set_servers( $args['servers'] ?? array() );
        $this->_debug = $args['debug'] ?? false;
        $this->stats = array();
        $this->_compress_threshold = $args['compress_threshold'] ?? 0;
        $this->_persistent = $args['persistent'] ?? false;
        $this->_compress_enable = true;
        $this->_have_zlib = function_exists( 'gzcompress' );
 
        $this->_cache_sock = array();
        $this->_host_dead = array();
 
        $this->_timeout_seconds = 0;
        $this->_timeout_microseconds = $args['timeout'] ?? 500000;
 
        $this->_connect_timeout = $args['connect_timeout'] ?? 0.1;
        $this->_connect_attempts = 2;
 
        $this->_logger = $args['logger'] ?? new NullLogger();
    }
 
    // }}}
 
    public function serialize( $value ) {
        return serialize( $value );
    }
 
    public function unserialize( $value ) {
        return unserialize( $value );
    }
 
    // {{{ add()
 
    public function add( $key, $val, $exp = 0 ) {
        return $this->_set( 'add', $key, $val, $exp );
    }
 
    // }}}
    // {{{ decr()
 
    public function decr( $key, $amt = 1 ) {
        return $this->_incrdecr( 'decr', $key, $amt );
    }
 
    // }}}
    // {{{ delete()
 
    public function delete( $key, $time = 0 ) {
        if ( !$this->_active ) {
            return false;
        }
 
        $sock = $this->get_sock( $key );
        if ( !$sock ) {
            return false;
        }
 
        $key = is_array( $key ) ? $key[1] : $key;
 
        if ( isset( $this->stats['delete'] ) ) {
            $this->stats['delete']++;
        } else {
            $this->stats['delete'] = 1;
        }
        $cmd = "delete $key $time\r\n";
        if ( !$this->_fwrite( $sock, $cmd ) ) {
            return false;
        }
        $res = $this->_fgets( $sock );
 
        if ( $this->_debug ) {
            $this->_debugprint( sprintf( "MemCache: delete %s (%s)", $key, $res ) );
        }
 
        if ( $res == "DELETED" || $res == "NOT_FOUND" ) {
            return true;
        }
 
        return false;
    }
 
    public function touch( $key, $time = 0 ) {
        if ( !$this->_active ) {
            return false;
        }
 
        $sock = $this->get_sock( $key );
        if ( !$sock ) {
            return false;
        }
 
        $key = is_array( $key ) ? $key[1] : $key;
 
        if ( isset( $this->stats['touch'] ) ) {
            $this->stats['touch']++;
        } else {
            $this->stats['touch'] = 1;
        }
        $cmd = "touch $key $time\r\n";
        if ( !$this->_fwrite( $sock, $cmd ) ) {
            return false;
        }
        $res = $this->_fgets( $sock );
 
        if ( $this->_debug ) {
            $this->_debugprint( sprintf( "MemCache: touch %s (%s)", $key, $res ) );
        }
 
        if ( $res == "TOUCHED" ) {
            return true;
        }
 
        return false;
    }
 
    // }}}
    // {{{ disconnect_all()
 
    public function disconnect_all() {
        foreach ( $this->_cache_sock as $sock ) {
            fclose( $sock );
        }
 
        $this->_cache_sock = array();
    }
 
    // }}}
    // {{{ enable_compress()
 
    public function enable_compress( $enable ) {
        $this->_compress_enable = $enable;
    }
 
    // }}}
    // {{{ forget_dead_hosts()
 
    public function forget_dead_hosts() {
        $this->_host_dead = array();
    }
 
    // }}}
    // {{{ get()
 
    public function get( $key, &$casToken = null ) {
        $getToken = ( func_num_args() >= 2 );
 
        if ( $this->_debug ) {
            $this->_debugprint( "get($key)" );
        }
 
        if ( !is_array( $key ) && strval( $key ) === '' ) {
            $this->_debugprint( "Skipping key which equals to an empty string" );
            return false;
        }
 
        if ( !$this->_active ) {
            return false;
        }
 
        $sock = $this->get_sock( $key );
 
        if ( !$sock ) {
            return false;
        }
 
        $key = is_array( $key ) ? $key[1] : $key;
        if ( isset( $this->stats['get'] ) ) {
            $this->stats['get']++;
        } else {
            $this->stats['get'] = 1;
        }
 
        $cmd = $getToken ? "gets" : "get";
        $cmd .= " $key\r\n";
        if ( !$this->_fwrite( $sock, $cmd ) ) {
            return false;
        }
 
        $val = array();
        $this->_load_items( $sock, $val, $casToken );
 
        if ( $this->_debug ) {
            foreach ( $val as $k => $v ) {
                $this->_debugprint(
                    sprintf( "MemCache: sock %s got %s", $this->serialize( $sock ), $k ) );
            }
        }
 
        $value = false;
        if ( isset( $val[$key] ) ) {
            $value = $val[$key];
        }
        return $value;
    }
 
    // }}}
    // {{{ get_multi()
 
    public function get_multi( $keys ) {
        if ( !$this->_active ) {
            return array();
        }
 
        if ( isset( $this->stats['get_multi'] ) ) {
            $this->stats['get_multi']++;
        } else {
            $this->stats['get_multi'] = 1;
        }
        $sock_keys = array();
        $socks = array();
        foreach ( $keys as $key ) {
            $sock = $this->get_sock( $key );
            if ( !$sock ) {
                continue;
            }
            $key = is_array( $key ) ? $key[1] : $key;
            $sockValue = intval( $sock );
 
            if ( !isset( $sock_keys[$sockValue] ) ) {
                $sock_keys[$sockValue] = array();
                $socks[] = $sock;
            }
            $sock_keys[$sockValue][] = $key;
        }
 
        $gather = array();
        // Send out the requests
        foreach ( $socks as $sock ) {
            $cmd = 'get';
            foreach ( $sock_keys[intval( $sock )] as $key ) {
                $cmd .= ' ' . $key;
            }
            $cmd .= "\r\n";
 
            if ( $this->_fwrite( $sock, $cmd ) ) {
                $gather[] = $sock;
            }
        }
 
        // Parse responses
        $val = array();
        foreach ( $gather as $sock ) {
            $this->_load_items( $sock, $val );
        }
 
        if ( $this->_debug ) {
            foreach ( $val as $k => $v ) {
                $this->_debugprint( sprintf( "MemCache: got %s", $k ) );
            }
        }
 
        return $val;
    }
 
    // }}}
    // {{{ incr()
 
    public function incr( $key, $amt = 1 ) {
        return $this->_incrdecr( 'incr', $key, $amt );
    }
 
    // }}}
    // {{{ replace()
 
    public function replace( $key, $value, $exp = 0 ) {
        return $this->_set( 'replace', $key, $value, $exp );
    }
 
    // }}}
    // {{{ run_command()
 
    public function run_command( $sock, $cmd ) {
        if ( !$sock ) {
            return array();
        }
 
        if ( !$this->_fwrite( $sock, $cmd ) ) {
            return array();
        }
 
        $ret = array();
        while ( true ) {
            $res = $this->_fgets( $sock );
            $ret[] = $res;
            if ( preg_match( '/^END/', $res ) ) {
                break;
            }
            if ( strlen( $res ) == 0 ) {
                break;
            }
        }
        return $ret;
    }
 
    // }}}
    // {{{ set()
 
    public function set( $key, $value, $exp = 0 ) {
        return $this->_set( 'set', $key, $value, $exp );
    }
 
    // }}}
    // {{{ cas()
 
    public function cas( $casToken, $key, $value, $exp = 0 ) {
        return $this->_set( 'cas', $key, $value, $exp, $casToken );
    }
 
    // }}}
    // {{{ set_compress_threshold()
 
    public function set_compress_threshold( $thresh ) {
        $this->_compress_threshold = $thresh;
    }
 
    // }}}
    // {{{ set_debug()
 
    public function set_debug( $dbg ) {
        $this->_debug = $dbg;
    }
 
    // }}}
    // {{{ set_servers()
 
    public function set_servers( $list ) {
        $this->_servers = $list;
        $this->_active = count( $list );
        $this->_buckets = null;
        $this->_bucketcount = 0;
 
        $this->_single_sock = null;
        if ( $this->_active == 1 ) {
            $this->_single_sock = $this->_servers[0];
        }
    }
 
    public function set_timeout( $seconds, $microseconds ) {
        $this->_timeout_seconds = $seconds;
        $this->_timeout_microseconds = $microseconds;
    }
 
    // }}}
    // }}}
    // {{{ private methods
    // {{{ _close_sock()
 
    function _close_sock( $sock ) {
        $host = array_search( $sock, $this->_cache_sock );
        fclose( $this->_cache_sock[$host] );
        unset( $this->_cache_sock[$host] );
    }
 
    // }}}
    // {{{ _connect_sock()
 
    function _connect_sock( &$sock, $host ) {
        $port = null;
        $hostAndPort = IPUtils::splitHostAndPort( $host );
        if ( $hostAndPort ) {
            $ip = $hostAndPort[0];
            if ( $hostAndPort[1] ) {
                $port = $hostAndPort[1];
            }
        } else {
            $ip = $host;
        }
        $sock = false;
        $timeout = $this->_connect_timeout;
        $errno = $errstr = null;
        for ( $i = 0; !$sock && $i < $this->_connect_attempts; $i++ ) {
            Wikimedia\suppressWarnings();
            if ( $this->_persistent == 1 ) {
                $sock = pfsockopen( $ip, $port, $errno, $errstr, $timeout );
            } else {
                $sock = fsockopen( $ip, $port, $errno, $errstr, $timeout );
            }
            Wikimedia\restoreWarnings();
        }
        if ( !$sock ) {
            $this->_error_log( "Error connecting to $host: $errstr" );
            $this->_dead_host( $host );
            return false;
        }
 
        // Initialise timeout
        stream_set_timeout( $sock, $this->_timeout_seconds, $this->_timeout_microseconds );
 
        // If the connection was persistent, flush the read buffer in case there
        // was a previous incomplete request on this connection
        if ( $this->_persistent ) {
            $this->_flush_read_buffer( $sock );
        }
        return true;
    }
 
    // }}}
    // {{{ _dead_sock()
 
    function _dead_sock( $sock ) {
        $host = array_search( $sock, $this->_cache_sock );
        $this->_dead_host( $host );
    }
 
    function _dead_host( $host ) {
        $hostAndPort = IPUtils::splitHostAndPort( $host );
        if ( $hostAndPort ) {
            $ip = $hostAndPort[0];
        } else {
            $ip = $host;
        }
        $this->_host_dead[$ip] = time() + 30 + intval( rand( 0, 10 ) );
        $this->_host_dead[$host] = $this->_host_dead[$ip];
        unset( $this->_cache_sock[$host] );
    }
 
    // }}}
    // {{{ get_sock()
 
    function get_sock( $key ) {
        if ( !$this->_active ) {
            return false;
        }
 
        if ( $this->_single_sock !== null ) {
            return $this->sock_to_host( $this->_single_sock );
        }
 
        $hv = is_array( $key ) ? intval( $key[0] ) : $this->_hashfunc( $key );
        if ( $this->_buckets === null ) {
            $bu = array();
            foreach ( $this->_servers as $v ) {
                if ( is_array( $v ) ) {
                    for ( $i = 0; $i < $v[1]; $i++ ) {
                        $bu[] = $v[0];
                    }
                } else {
                    $bu[] = $v;
                }
            }
            $this->_buckets = $bu;
            $this->_bucketcount = count( $bu );
        }
 
        $realkey = is_array( $key ) ? $key[1] : $key;
        for ( $tries = 0; $tries < 20; $tries++ ) {
            $host = $this->_buckets[$hv % $this->_bucketcount];
            $sock = $this->sock_to_host( $host );
            if ( $sock ) {
                return $sock;
            }
            $hv = $this->_hashfunc( $hv . $realkey );
        }
 
        return false;
    }
 
    // }}}
    // {{{ _hashfunc()
 
    function _hashfunc( $key ) {
        # Hash function must be in [0,0x7ffffff]
        # We take the first 31 bits of the MD5 hash, which unlike the hash
        # function used in a previous version of this client, works
        return hexdec( substr( md5( $key ), 0, 8 ) ) & 0x7fffffff;
    }
 
    // }}}
    // {{{ _incrdecr()
 
    function _incrdecr( $cmd, $key, $amt = 1 ) {
        if ( !$this->_active ) {
            return null;
        }
 
        $sock = $this->get_sock( $key );
        if ( !$sock ) {
            return null;
        }
 
        $key = is_array( $key ) ? $key[1] : $key;
        if ( isset( $this->stats[$cmd] ) ) {
            $this->stats[$cmd]++;
        } else {
            $this->stats[$cmd] = 1;
        }
        if ( !$this->_fwrite( $sock, "$cmd $key $amt\r\n" ) ) {
            return null;
        }
 
        $line = $this->_fgets( $sock );
        $match = array();
        if ( !preg_match( '/^(\d+)/', $line, $match ) ) {
            return null;
        }
        return $match[1];
    }
 
    // }}}
    // {{{ _load_items()
 
    function _load_items( $sock, &$ret, &$casToken = null ) {
        $results = array();
 
        while ( 1 ) {
            $decl = $this->_fgets( $sock );
 
            if ( $decl === false ) {
                /*
                 * If nothing can be read, something is wrong because we know exactly when
                 * to stop reading (right after "END") and we return right after that.
                 */
                return false;
            } elseif ( preg_match( '/^VALUE (\S+) (\d+) (\d+)(?: (\d+))?$/', $decl, $match ) ) {
                /*
                 * Read all data returned. This can be either one or multiple values.
                 * Save all that data (in an array) to be processed later: we'll first
                 * want to continue reading until "END" before doing anything else,
                 * to make sure that we don't leave our client in a state where it's
                 * output is not yet fully read.
                 */
                $results[] = array(
                    $match[1], // rkey
                    $match[2], // flags
                    $match[3], // len
                    $match[4] ?? null, // casToken (appears with "gets" but not "get")
                    $this->_fread( $sock, $match[3] + 2 ), // data
                );
            } elseif ( $decl == "END" ) {
                if ( count( $results ) == 0 ) {
                    return false;
                }
 
                foreach ( $results as [ $rkey, $flags, /* length */, $casToken, $data ] ) {
                    if ( $data === false || substr( $data, -2 ) !== "\r\n" ) {
                        $this->_handle_error( $sock,
                            'line ending missing from data block from $1' );
                        return false;
                    }
                    $data = substr( $data, 0, -2 );
                    $ret[$rkey] = $data;
 
                    if ( $this->_have_zlib && $flags & self::COMPRESSED ) {
                        $ret[$rkey] = gzuncompress( $ret[$rkey] );
                    }
 
                    /*
                     * This unserialize is the exact reason that we only want to
                     * process data after having read until "END" (instead of doing
                     * this right away): "unserialize" can trigger outside code:
                     * in the event that $ret[$rkey] is a serialized object,
                     * unserializing it will trigger __wakeup() if present. If that
                     * function attempted to read from memcached (while we did not
                     * yet read "END"), these 2 calls would collide.
                     */
                    if ( $flags & self::SERIALIZED ) {
                        $ret[$rkey] = $this->unserialize( $ret[$rkey] );
                    } elseif ( $flags & self::INTVAL ) {
                        $ret[$rkey] = intval( $ret[$rkey] );
                    }
                }
 
                return true;
            } else {
                $this->_handle_error( $sock, 'Error parsing response from $1' );
                return false;
            }
        }
    }
 
    // }}}
    // {{{ _set()
 
    function _set( $cmd, $key, $val, $exp, $casToken = null ) {
        if ( !$this->_active ) {
            return false;
        }
 
        $sock = $this->get_sock( $key );
        if ( !$sock ) {
            return false;
        }
 
        if ( isset( $this->stats[$cmd] ) ) {
            $this->stats[$cmd]++;
        } else {
            $this->stats[$cmd] = 1;
        }
 
        $flags = 0;
 
        if ( is_int( $val ) ) {
            $flags |= self::INTVAL;
        } elseif ( !is_scalar( $val ) ) {
            $val = $this->serialize( $val );
            $flags |= self::SERIALIZED;
            if ( $this->_debug ) {
                $this->_debugprint( sprintf( "client: serializing data as it is not scalar" ) );
            }
        }
 
        $len = strlen( $val );
 
        if ( $this->_have_zlib && $this->_compress_enable
            && $this->_compress_threshold && $len >= $this->_compress_threshold
        ) {
            $c_val = gzcompress( $val, 9 );
            $c_len = strlen( $c_val );
 
            if ( $c_len < $len * ( 1 - self::COMPRESSION_SAVINGS ) ) {
                if ( $this->_debug ) {
                    $this->_debugprint( sprintf( "client: compressing data; was %d bytes is now %d bytes", $len, $c_len ) );
                }
                $val = $c_val;
                $len = $c_len;
                $flags |= self::COMPRESSED;
            }
        }
 
        $command = "$cmd $key $flags $exp $len";
        if ( $casToken ) {
            $command .= " $casToken";
        }
 
        if ( !$this->_fwrite( $sock, "$command\r\n$val\r\n" ) ) {
            return false;
        }
 
        $line = $this->_fgets( $sock );
 
        if ( $this->_debug ) {
            $this->_debugprint( sprintf( "%s %s (%s)", $cmd, $key, $line ) );
        }
        if ( $line === "STORED" ) {
            return true;
        } elseif ( $line === "NOT_STORED" && $cmd === "set" ) {
            // "Not stored" is always used as the mcrouter response with AllAsyncRoute
            return true;
        }
 
        return false;
    }
 
    // }}}
    // {{{ sock_to_host()
 
    function sock_to_host( $host ) {
        if ( isset( $this->_cache_sock[$host] ) ) {
            return $this->_cache_sock[$host];
        }
 
        $sock = null;
        $now = time();
        $hostAndPort = IPUtils::splitHostAndPort( $host );
        if ( $hostAndPort ) {
            $ip = $hostAndPort[0];
        } else {
            $ip = $host;
        }
        if ( isset( $this->_host_dead[$host] ) && $this->_host_dead[$host] > $now ||
            isset( $this->_host_dead[$ip] ) && $this->_host_dead[$ip] > $now
        ) {
            return null;
        }
 
        if ( !$this->_connect_sock( $sock, $host ) ) {
            return null;
        }
 
        // Do not buffer writes
        stream_set_write_buffer( $sock, 0 );
 
        $this->_cache_sock[$host] = $sock;
 
        return $this->_cache_sock[$host];
    }
 
    function _debugprint( $text ) {
        $this->_logger->debug( $text );
    }
 
    function _error_log( $text ) {
        $this->_logger->error( "Memcached error: $text" );
    }
 
    function _fwrite( $sock, $buf ) {
        $bytesWritten = 0;
        $bufSize = strlen( $buf );
        while ( $bytesWritten < $bufSize ) {
            $result = fwrite( $sock, $buf );
            $data = stream_get_meta_data( $sock );
            if ( $data['timed_out'] ) {
                $this->_handle_error( $sock, 'timeout writing to $1' );
                return false;
            }
            // Contrary to the documentation, fwrite() returns zero on error in PHP 5.3.
            if ( $result === false || $result === 0 ) {
                $this->_handle_error( $sock, 'error writing to $1' );
                return false;
            }
            $bytesWritten += $result;
        }
 
        return true;
    }
 
    function _handle_error( $sock, $msg ) {
        $peer = stream_socket_get_name( $sock, true  );
        if ( strval( $peer ) === '' ) {
            $peer = array_search( $sock, $this->_cache_sock );
            if ( $peer === false ) {
                $peer = '[unknown host]';
            }
        }
        $msg = str_replace( '$1', $peer, $msg );
        $this->_error_log( "$msg" );
        $this->_dead_sock( $sock );
    }
 
    function _fread( $sock, $len ) {
        $buf = '';
        while ( $len > 0 ) {
            $result = fread( $sock, $len );
            $data = stream_get_meta_data( $sock );
            if ( $data['timed_out'] ) {
                $this->_handle_error( $sock, 'timeout reading from $1' );
                return false;
            }
            if ( $result === false ) {
                $this->_handle_error( $sock, 'error reading buffer from $1' );
                return false;
            }
            if ( $result === '' ) {
                // This will happen if the remote end of the socket is shut down
                $this->_handle_error( $sock, 'unexpected end of file reading from $1' );
                return false;
            }
            $len -= strlen( $result );
            $buf .= $result;
        }
        return $buf;
    }
 
    function _fgets( $sock ) {
        $result = fgets( $sock );
        // fgets() may return a partial line if there is a select timeout after
        // a successful recv(), so we have to check for a timeout even if we
        // got a string response.
        $data = stream_get_meta_data( $sock );
        if ( $data['timed_out'] ) {
            $this->_handle_error( $sock, 'timeout reading line from $1' );
            return false;
        }
        if ( $result === false ) {
            $this->_handle_error( $sock, 'error reading line from $1' );
            return false;
        }
        if ( substr( $result, -2 ) === "\r\n" ) {
            $result = substr( $result, 0, -2 );
        } elseif ( substr( $result, -1 ) === "\n" ) {
            $result = substr( $result, 0, -1 );
        } else {
            $this->_handle_error( $sock, 'line ending missing in response from $1' );
            return false;
        }
        return $result;
    }
 
    function _flush_read_buffer( $f ) {
        if ( !$f ) {
            return;
        }
        $r = array( $f );
        $w = null;
        $e = null;
        $n = stream_select( $r, $w, $e, 0, 0 );
        while ( $n == 1 && !feof( $f ) ) {
            fread( $f, 1024 );
            $r = array( $f );
            $w = null;
            $e = null;
            $n = stream_select( $r, $w, $e, 0, 0 );
        }
    }
 
    // }}}
    // }}}
    // }}}
}
 
// }}}