extra/memcached/memcached.factor

   1 ! Copyright (C) 2010 John Benediktsson
   2 ! See http://factorcode.org/license.txt for BSD license
   3
   4 USING: accessors arrays assocs byte-arrays combinators fry
   5 io io.encodings.binary io.sockets kernel make math math.parser
   6 namespaces pack random sequences strings ;
   7
   8 IN: memcached
   9
  10 ! TODO:
  11 ! - quiet commands
  12 ! - CAS
  13 ! - expirations
  14 ! - initial-value for incr/decr
  15
  16
  17 SYMBOL: memcached-server
  18 "127.0.0.1" 11211 <inet> memcached-server set-global
  19
  20 : with-memcached ( quot -- )
  21     memcached-server get-global
  22     binary [ call ] with-client ; inline
  23
  24 ERROR: key-not-found ;
  25 ERROR: key-exists ;
  26 ERROR: value-too-large ;
  27 ERROR: invalid-arguments ;
  28 ERROR: item-not-stored ;
  29 ERROR: value-not-numeric ;
  30 ERROR: unknown-command ;
  31 ERROR: out-of-memory ;
  32
  33 <PRIVATE
  34
  35 ! Commands
  36 CONSTANT: GET      0x00
  37 CONSTANT: SET      0x01
  38 CONSTANT: ADD      0x02
  39 CONSTANT: REPLACE  0x03
  40 CONSTANT: DELETE   0x04
  41 CONSTANT: INCR     0x05
  42 CONSTANT: DECR     0x06
  43 CONSTANT: QUIT     0x07
  44 CONSTANT: FLUSH    0x08
  45 CONSTANT: GETQ     0x09
  46 CONSTANT: NOOP     0x0A
  47 CONSTANT: VERSION  0x0B
  48 CONSTANT: GETK     0x0C
  49 CONSTANT: GETKQ    0x0D
  50 CONSTANT: APPEND   0x0E
  51 CONSTANT: PREPEND  0x0F
  52 CONSTANT: STAT     0x10
  53 CONSTANT: SETQ     0x11
  54 CONSTANT: ADDQ     0x12
  55 CONSTANT: REPLACEQ 0x13
  56 CONSTANT: DELETEQ  0x14
  57 CONSTANT: INCRQ    0x15
  58 CONSTANT: DECRQ    0x16
  59 CONSTANT: QUITQ    0x17
  60 CONSTANT: FLUSHQ   0x18
  61 CONSTANT: APPENDQ  0x19
  62 CONSTANT: PREPENDQ 0x1A
  63
  64 ! Errors
  65 CONSTANT: NOT_FOUND    0x01
  66 CONSTANT: EXISTS       0x02
  67 CONSTANT: TOO_LARGE    0x03
  68 CONSTANT: INVALID_ARGS 0x04
  69 CONSTANT: NOT_STORED   0x05
  70 CONSTANT: NOT_NUMERIC  0x06
  71 CONSTANT: UNKNOWN_CMD  0x81
  72 CONSTANT: MEMORY       0x82
  73
  74 TUPLE: request cmd key val extra opaque cas ;
  75
  76 : <request> ( cmd -- request )
  77     "" "" "" random-32 0 \ request boa ;
  78
  79 : send-header ( request -- )
  80     {
  81         [ cmd>> ]
  82         [ key>> length ]
  83         [ extra>> length ]
  84         [
  85             [ key>> length ]
  86             [ extra>> length ]
  87             [ val>> length ] tri + +
  88         ]
  89         [ opaque>> ]
  90         [ cas>> ]
  91     } cleave
  92     ! magic, opcode, keylen, extralen, datatype, status,
  93     ! bodylen, opaque, cas [ big-endian ]
  94     '[ 0x80 _ _ _ 0 0 _ _ _ ] "CCSCCSIIQ" pack-be write ;
  95
  96 : (send) ( str -- )
  97     [ >byte-array write ] unless-empty ;
  98
  99 : send-request ( request -- )
 100     {
 101         [ send-header    ]
 102         [ extra>> (send) ]
 103         [ key>>   (send) ]
 104         [ val>>   (send) ]
 105     } cleave flush ;
 106
 107 : read-header ( -- header )
 108     "CCSCCSIIQ" [ packed-length read ] [ unpack-be ] bi ;
 109
 110 : check-magic ( header -- )
 111     first 0x81 = [ "bad magic" throw ] unless ;
 112
 113 : check-status ( header -- )
 114     [ 5 ] dip nth {
 115         { NOT_FOUND    [ key-not-found     ] }
 116         { EXISTS       [ key-exists        ] }
 117         { TOO_LARGE    [ value-too-large   ] }
 118         { INVALID_ARGS [ invalid-arguments ] }
 119         { NOT_STORED   [ item-not-stored   ] }
 120         { NOT_NUMERIC  [ value-not-numeric ] }
 121         { UNKNOWN_CMD  [ unknown-command   ] }
 122         { MEMORY       [ out-of-memory     ] }
 123         [ drop ]
 124     } case ;
 125
 126 : check-opaque ( opaque header -- ? )
 127     [ 7 ] dip nth = ;
 128
 129 : (read) ( n -- str )
 130     dup 0 > [ read >string ] [ drop "" ] if ;
 131
 132 : read-key ( header -- key )
 133     [ 2 ] dip nth (read) ;
 134
 135 : read-val ( header -- val )
 136     [ [ 6 ] dip nth ] [ [ 2 ] dip nth ] bi - (read) ;
 137
 138 : read-body ( header -- val key )
 139     {
 140         [ check-magic  ]
 141         [ check-status ]
 142         [ read-key     ]
 143         [ read-val     ]
 144     } cleave swap ;
 145
 146 : read-response ( -- val key )
 147     read-header read-body ;
 148
 149 : submit ( request -- response )
 150     send-request read-response drop ;
 151
 152 : (cmd) ( key cmd -- request )
 153     <request> swap >>key ;
 154
 155 : (incr/decr) ( amt key cmd -- response )
 156     (cmd) swap '[ _ 0 0 ] "QQI" pack-be >>extra ! amt init exp
 157     submit "Q" unpack-be first ;
 158
 159 : (mutate) ( val key cmd -- )
 160     (cmd) swap >>val { 0 0 } "II" pack-be >>extra ! flags exp
 161     submit drop ;
 162
 163 : (cat) ( val key cmd -- )
 164     (cmd) swap >>val submit drop ;
 165
 166 PRIVATE>
 167
 168 : m/version ( -- version ) VERSION <request> submit ;
 169
 170 : m/noop ( -- ) NOOP <request> submit drop ;
 171
 172 : m/incr-val ( amt key -- val ) INCR (incr/decr) ;
 173
 174 : m/incr ( key -- val ) 1 swap m/incr-val ;
 175
 176 : m/decr-val ( amt key -- val ) DECR (incr/decr) ;
 177
 178 : m/decr ( key -- val ) 1 swap m/decr-val ;
 179
 180 : m/get ( key -- val ) GET (cmd) submit 4 tail ;
 181
 182 : m/getq ( opaque key -- )
 183     GETQ (cmd) swap >>opaque send-request ;
 184
 185 : m/getseq ( keys -- vals )
 186     [ H{ } clone ] dip
 187     [ <enum> [ m/getq ] assoc-each ]
 188     [ length 10 + NOOP <request> swap >>opaque send-request ]
 189     [
 190         <enum> [
 191             assoc-size 10 + '[
 192                 _ read-header [ check-opaque ] keep swap
 193             ]
 194         ] [
 195             '[
 196                 [ read-body drop 4 tail ]
 197                 [ [ 7 ] dip nth _ at ]
 198                 bi pick set-at
 199             ]
 200         ] bi until drop
 201     ] tri ;
 202
 203 : m/set ( val key -- ) SET (mutate) ;
 204
 205 : m/add ( val key -- ) ADD (mutate) ;
 206
 207 : m/replace ( val key -- ) REPLACE (mutate) ;
 208
 209 : m/delete ( key -- ) DELETE (cmd) submit drop ;
 210
 211 : m/append ( val key -- ) APPEND (cat) ;
 212
 213 : m/prepend ( val key -- ) PREPEND (cat) ;
 214
 215 : m/flush-later ( seconds -- )
 216     FLUSH <request> swap 1array "I" pack-be >>extra ! timebomb
 217     submit drop ;
 218
 219 : m/flush ( -- ) 0 m/flush-later ;
 220
 221 : m/stats ( -- stats )
 222     STAT <request> send-request
 223     [ read-response dup length 0 > ]
 224     [ swap 2array ] produce 2nip ;
 225
 226 : m/quit ( -- ) QUIT <request> submit drop ;
 227
 228