Paste: memcached

! Copyright (C) 2010 John Benediktsson
! See http://factorcode.org/license.txt for BSD license

USING: accessors arrays byte-arrays combinators fry io
io.encodings.binary io.sockets kernel make math
namespaces pack random sequences strings ;

IN: memcached

! Commands
CONSTANT: GET 0
CONSTANT: SET 1
CONSTANT: ADD 2
CONSTANT: REPLACE 3
CONSTANT: DELETE 4
CONSTANT: INCR 5
CONSTANT: DECR 6
CONSTANT: QUIT 7
CONSTANT: FLUSH 8
CONSTANT: GETQ 9
CONSTANT: NOOP 10
CONSTANT: VERSION 11
CONSTANT: STAT HEX: 10
CONSTANT: APPEND HEX: 0e
CONSTANT: PREPEND HEX: 0f

! Errors
CONSTANT: UNKNOWN_CMD HEX: 81
CONSTANT: NOT_FOUND HEX: 1
CONSTANT: EXISTS HEX: 2

SYMBOL: memcached-server
"127.0.0.1" 11211 <inet> memcached-server set-global

: with-memcached ( quot -- )
    memcached-server get-global
    binary [ call ] with-client ; inline

<PRIVATE

TUPLE: request cmd key val extra opaque cas ;

: <request> ( cmd -- request )
    "" "" "" random-32 0 \ request boa ;

: send-header ( request -- )
    {
        [ cmd>> ]
        [ key>> length ]
        [ extra>> length ]
        [
            [ key>> length ]
            [ extra>> length ]
            [ val>> length ] tri + +
        ]
        [ opaque>> ]
        [ cas>> ]
    } cleave
    '[ HEX: 80 _ _ _ 0 0 _ _ _ ] "CCSCCSIIQ" pack-be write ;

: (send) ( str -- )
    >byte-array dup length 0 > [ write ] [ drop ] if ;

: send-request ( request -- )
    {
        [ send-header    ]
        [ extra>> (send) ]
        [ key>>   (send) ]
        [ val>>   (send) ]
    } cleave flush ;


: read-header ( -- header )
    "CCSCCSIIQ" [ packed-length read ] [ unpack-be ] bi ;

: check-magic ( header -- )
    first HEX: 81 = [ "bad magic" throw ] unless ;

: check-status ( header -- )
    [ 5 ] dip nth {
        { UNKNOWN_CMD [ "unknown command" throw ] }
        { NOT_FOUND   [ "not found" throw ] }
        { EXISTS      [ "exists" throw ] }
        [ drop ]
    } case ;

: (read) ( n -- str )
    dup 0 > [ read >string ] [ drop "" ] if ;

: read-key ( header -- key )
    [ 2 ] dip nth (read) ;

: read-val ( header -- body )
    [ [ 6 ] dip nth ] [ [ 2 ] dip nth ] bi - (read) ;

: read-response ( -- val key )
    read-header {
        [ check-magic  ]
        [ check-status ]
        [ read-key     ]
        [ read-val     ]
    } cleave swap ;

: read-responses ( -- responses )
    [
        [ read-response dup length 0 > ]
        [ swap 2array , ] while drop
    ] { } make nip ;


: single ( request -- response )
    send-request read-response drop ;

: multi ( request -- responses )
    send-request read-responses ;


: (cmd) ( key cmd -- request )
    <request> swap >>key ;

: (incr/decr) ( amt key cmd -- response )
    (cmd)
    swap '[ _ 0 0 ] "QQI" pack-be >>extra ! amt init exp
    single "Q" unpack-be first ;

: (mutate) ( val key cmd -- )
    (cmd) swap >>val
    { 0 0 } "II" pack-be >>extra ! flags, expiration
    single drop ;

: (cat) ( val key cmd -- )
    (cmd) swap >>val single drop ;

PRIVATE>

: m/version ( -- version ) VERSION <request> single ;

: m/noop ( -- ) NOOP <request> single drop ;

: m/stats ( -- stats ) STAT <request> multi ;

: m/incr-val ( amt key -- result ) INCR (incr/decr) ;

: m/incr ( key -- result ) 1 swap m/incr-val ;

: m/decr-val ( amt key -- result ) DECR (incr/decr) ;

: m/decr ( key -- result ) 1 swap m/decr-val ;

: m/get ( key -- result ) GET (cmd) single 4 tail ;

: m/set ( val key -- ) SET (mutate) ;

: m/add ( val key -- ) ADD (mutate) ;

: m/replace ( val key -- ) REPLACE (mutate) ;

: m/delete ( key -- ) DELETE (cmd) single drop ;

: m/append ( val key -- ) APPEND (cat) ;

: m/prepend ( val key -- ) PREPEND (cat) ;

: m/flush-later ( time -- )
    FLUSH <request>
    swap 1array "I" pack-be >>extra ! timebomb
    single drop ;

: m/flush ( -- ) 0 m/flush-later ;