001/*
002 * (C) Copyright 2017 Nuxeo (http://nuxeo.com/) and others.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *     http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 *
016 * Contributors:
017 *     Florent Guillaume
018 */
019package org.nuxeo.runtime.kv;
020
021import java.lang.reflect.Field;
022import java.util.Arrays;
023import java.util.Objects;
024import java.util.concurrent.TimeUnit;
025import java.util.concurrent.locks.Lock;
026import java.util.stream.Stream;
027
028import net.jodah.expiringmap.ExpiringMap;
029
030/**
031 * Memory-based implementation of a Key/Value store.
032 *
033 * @since 9.1
034 */
035public class MemKeyValueStore extends AbstractKeyValueStoreProvider {
036
037    protected final ExpiringMap<String, byte[]> map;
038
039    protected final Lock writeLock;
040
041    protected String name;
042
043    public MemKeyValueStore() {
044        map = ExpiringMap.builder().expiration(Integer.MAX_VALUE, TimeUnit.DAYS).variableExpiration().build();
045        try {
046            Field field = map.getClass().getDeclaredField("writeLock");
047            field.setAccessible(true);
048            writeLock = (Lock) field.get(map);
049        } catch (ReflectiveOperationException | SecurityException e) {
050            throw new RuntimeException(e);
051        }
052    }
053
054    @Override
055    public void initialize(KeyValueStoreDescriptor descriptor) {
056        this.name = descriptor.name;
057    }
058
059    @Override
060    public Stream<String> keyStream() {
061        return map.keySet().stream();
062    }
063
064    @Override
065    public void close() {
066    }
067
068    @Override
069    public void clear() {
070        map.clear();
071    }
072
073    protected static byte[] clone(byte[] value) {
074        return value == null ? null : value.clone();
075    }
076
077    @Override
078    public void put(String key, byte[] value, long ttl) {
079        Objects.requireNonNull(key);
080        value = clone(value);
081        if (value == null) {
082            map.remove(key);
083        } else if (ttl == 0) {
084            map.put(key, value);
085        } else {
086            map.put(key, value, ttl, TimeUnit.SECONDS);
087        }
088    }
089
090    @Override
091    public byte[] get(String key) {
092        Objects.requireNonNull(key);
093        byte[] value = map.get(key);
094        return clone(value);
095    }
096
097    @Override
098    public boolean setTTL(String key, long ttl) {
099        Objects.requireNonNull(key);
100        byte[] value = map.get(key);
101        if (value == null) {
102            return false;
103        }
104        doSetTTL(key, ttl);
105        return true;
106    }
107
108    protected void doSetTTL(String key, long ttl) {
109        if (ttl == 0) {
110            map.setExpiration(key, Integer.MAX_VALUE, TimeUnit.DAYS);
111        } else {
112            map.setExpiration(key, ttl, TimeUnit.SECONDS);
113        }
114    }
115
116    @Override
117    public boolean compareAndSet(String key, byte[] expected, byte[] value, long ttl) {
118        Objects.requireNonNull(key);
119        // clone is not needed if the comparison fails
120        // but we are optimistic and prefer to do the clone outside the lock
121        value = clone(value);
122        // we don't use ExpiringMap.replace because it deals with null differently
123        writeLock.lock();
124        try {
125            byte[] current = map.get(key);
126            boolean equal = Arrays.equals(expected, current);
127            if (equal) {
128                if (value == null) {
129                    map.remove(key);
130                } else {
131                    map.put(key, value);
132                    doSetTTL(key, ttl);
133                }
134            }
135            return equal;
136        } finally {
137            writeLock.unlock();
138        }
139    }
140
141    @Override
142    public String toString() {
143        return getClass().getSimpleName() + "(" + name + ")";
144    }
145
146}