001/**
002 * Copyright (C) 2006-2025 Talend Inc. - www.talend.com
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 */
016package org.talend.sdk.component.runtime.manager.service;
017
018import java.io.ObjectStreamException;
019import java.io.Serializable;
020import java.util.List;
021import java.util.Map.Entry;
022import java.util.Objects;
023import java.util.Optional;
024import java.util.concurrent.ConcurrentHashMap;
025import java.util.concurrent.ConcurrentMap;
026import java.util.concurrent.ScheduledExecutorService;
027import java.util.concurrent.ScheduledFuture;
028import java.util.concurrent.TimeUnit;
029import java.util.function.Function;
030import java.util.function.Predicate;
031import java.util.function.Supplier;
032import java.util.stream.Collectors;
033import java.util.stream.Stream;
034
035import javax.annotation.PreDestroy;
036
037import org.talend.sdk.component.api.configuration.Option;
038import org.talend.sdk.component.api.service.cache.LocalCache;
039import org.talend.sdk.component.api.service.configuration.Configuration;
040import org.talend.sdk.component.runtime.serialization.SerializableService;
041
042import lombok.Data;
043
044/**
045 * Implementation of LocalCache with in memory concurrent map.
046 */
047public class LocalCacheService implements LocalCache, Serializable {
048
049    /** plugin name for this cache */
050    private final String plugin;
051
052    private final Supplier<Long> timer;
053
054    private final ConcurrentMap<String, ElementImpl> cache = new ConcurrentHashMap<>();
055
056    @Configuration("talend.component.manager.services.cache.eviction")
057    private Supplier<CacheConfiguration> configuration;
058
059    // scheduler we use to evict tokens
060    private transient Supplier<ScheduledExecutorService> threadServiceGetter;
061
062    public LocalCacheService(final String plugin, final Supplier<Long> timer,
063            final Supplier<ScheduledExecutorService> threadServiceGetter) {
064        this.plugin = plugin;
065        this.timer = timer;
066        this.threadServiceGetter = threadServiceGetter;
067    }
068
069    /**
070     * Evict an object.
071     * 
072     * @param key key to evict.
073     */
074    @Override
075    public void evict(final String key) {
076        final String realKey = internalKey(key);
077
078        // use compute to be able to call release.
079        cache.compute(realKey, (String oldKey, ElementImpl oldElement) -> {
080            if (oldElement != null && oldElement.canBeEvict()) {
081                // ok to evict, so do release.
082                oldElement.release();
083                return null;
084            }
085            return oldElement;
086        });
087    }
088
089    @Override
090    public void evictIfValue(final String key, final Object expected) {
091        final String realKey = internalKey(key);
092
093        // use compute to be able to call release.
094        cache.compute(realKey, (String oldKey, ElementImpl oldElement) -> {
095            if (oldElement != null && (Objects.equals(oldElement.getValue(), expected) || oldElement.canBeEvict())) {
096                // ok to evit, so do release.
097                oldElement.release();
098                return null;
099            }
100            return oldElement;
101        });
102
103    }
104
105    @Override
106    public <T> T computeIfAbsent(final Class<T> expectedClass, final String key, final Predicate<Element> toRemove,
107            final long timeoutMs, final Supplier<T> value) {
108
109        final Integer maxSize = this.getConfigValue(CacheConfiguration::getDefaultMaxSize, -1);
110        if (maxSize > 0 && this.cache.size() >= maxSize) {
111            this.clean(); // clean before add one element.
112            if (this.cache.size() >= maxSize) {
113                synchronized (this.cache) {
114                    while (this.cache.size() >= maxSize) {
115                        final String keyToRemove = this.cache.keySet().iterator().next();
116                        this.cache.remove(keyToRemove);
117                    }
118                }
119            }
120        }
121
122        final ScheduledFuture<?> task = timeoutMs > 0 ? this.evictionTask(key, timeoutMs) : null;
123
124        final long endOfValidity = this.calcEndOfValidity(timeoutMs);
125        final ElementImpl element =
126                this.addToMap(key, () -> new ElementImpl(value, toRemove, endOfValidity, task, this.timer));
127
128        return element.getValue(expectedClass);
129    }
130
131    @Override
132    public <T> T computeIfAbsent(final Class<T> expectedClass, final String key, final Predicate<Element> toRemove,
133            final Supplier<T> value) {
134
135        final long timeout = this.getConfigValue(CacheConfiguration::getDefaultEvictionTimeout, -1L);
136        return this.computeIfAbsent(expectedClass, key, toRemove, timeout, value);
137    }
138
139    @Override
140    public <T> T computeIfAbsent(final Class<T> expectedClass, final String key, final long timeoutMs,
141            final Supplier<T> value) {
142        return this.computeIfAbsent(expectedClass, key, null, timeoutMs, value);
143    }
144
145    private ElementImpl addToMap(final String key, final Supplier<ElementImpl> builder) {
146        final String internalKey = internalKey(key);
147        return cache.compute(internalKey, (String k, ElementImpl old) -> //
148        old == null || old.mustBeRemoved() ? builder.get() : old);
149    }
150
151    @Override
152    public <T> T computeIfAbsent(final Class<T> expectedClass, final String key, final Supplier<T> value) {
153        final long timeOut = this.getConfigValue(CacheConfiguration::getDefaultEvictionTimeout, -1L);
154        return computeIfAbsent(expectedClass, key, null, timeOut, value);
155    }
156
157    @PreDestroy
158    public void release() {
159        this.cache.forEach((String k, ElementImpl e) -> e.release());
160        this.cache.clear();
161    }
162
163    private long calcEndOfValidity(final long timeoutMs) {
164        return timeoutMs > 0 ? this.timer.get() + timeoutMs : -1;
165    }
166
167    private String internalKey(final String key) {
168        return plugin + '@' + key;
169    }
170
171    public void clean() {
172        Stream<Entry<String, ElementImpl>> elements = //
173                this.cache
174                        .entrySet() //
175                        .stream() //
176                        .filter(e -> e.getValue().mustBeRemoved());
177
178        final int maxEviction = this.getConfigValue(CacheConfiguration::getMaxDeletionPerEvictionRun, -1);
179        if (maxEviction > 0) {
180            elements = elements.limit(maxEviction);
181        }
182        final List<String> removableElements = elements.map(Entry::getKey).collect(Collectors.toList());// materialize
183                                                                                                        // before
184                                                                                                        // actually
185                                                                                                        // removing it
186        removableElements.forEach(this.cache::remove);
187    }
188
189    private ScheduledExecutorService getThreadService() {
190        return this.threadServiceGetter.get();
191    }
192
193    /**
194     * Schedule an eviction task for a key.
195     * 
196     * @param key : key to evict from cache.
197     * @param delayMillis : delay in millis before triggered.
198     * @return result of task.
199     */
200    private ScheduledFuture<?> evictionTask(final String key, final long delayMillis) {
201        return this
202                .getThreadService()
203                .schedule(() -> this.evict(key), //
204                        delayMillis, //
205                        TimeUnit.MILLISECONDS); //
206    }
207
208    private <T> T getConfigValue(final Function<CacheConfiguration, T> getter, final T defaultValue) {
209        return Optional
210                .ofNullable(this.getConfig()) //
211                .map(getter) //
212                .orElse(defaultValue);
213    }
214
215    private CacheConfiguration getConfig() {
216        return this.configuration != null ? this.configuration.get() : null;
217    }
218
219    /**
220     * Cache configuration.
221     */
222    @Data
223    public static class CacheConfiguration implements Serializable {
224
225        @Option
226        private long defaultEvictionTimeout;
227
228        @Option
229        private int maxDeletionPerEvictionRun;
230
231        @Option
232        private int defaultMaxSize;
233    }
234
235    /**
236     * Wrapper for each cached object.
237     */
238    private static class ElementImpl implements Element {
239
240        /** cached object */
241        private final Object value;
242
243        /** function, if exists, that authorize to remove object. */
244        private final Predicate<Element> canBeRemoved;
245
246        /** give time object can be release (infinity if < 0) */
247        private final long endOfValidity;
248
249        /** scheduled task to remove object if nedeed (to cancel if removed before) */
250        private final ScheduledFuture<?> removedTask;
251
252        private final Supplier<Long> serviceTimer;
253
254        public <T> ElementImpl(final Supplier<T> value, final Predicate<Element> canBeRemoved, final long endOfValidity,
255                final ScheduledFuture<?> removedTask, final Supplier<Long> timer) {
256            this.value = value.get();
257            this.canBeRemoved = canBeRemoved;
258            this.endOfValidity = endOfValidity;
259            this.removedTask = removedTask;
260            this.serviceTimer = timer;
261        }
262
263        @Override
264        public <T> T getValue(final Class<T> expectedType) {
265            if (this.value != null && !expectedType.isInstance(this.value)) {
266                throw new ClassCastException(
267                        this.value.getClass().getName() + " cannot be cast to " + expectedType.getName());
268            }
269            return expectedType.cast(this.value);
270        }
271
272        @Override
273        public long getLastValidityTimestamp() {
274            return this.endOfValidity;
275        }
276
277        public boolean mustBeRemoved() {
278            return (this.endOfValidity > 0 && this.endOfValidity <= this.serviceTimer.get()) // time out passed
279                    && this.canBeEvict(); // or function indicate to remove.
280        }
281
282        public boolean canBeEvict() {
283            return this.canBeRemoved == null || this.canBeRemoved.test(this);
284        }
285
286        /**
287         * Release this object because removed.
288         */
289        public synchronized void release() {
290            if (this.removedTask != null) {
291                this.removedTask.cancel(false);
292            }
293        }
294
295        @Override
296        public boolean equals(final Object o) { // consider only value
297            if (this == o) {
298                return true;
299            }
300            if (o == null || getClass() != o.getClass()) {
301                return false;
302            }
303            return Objects.equals(ElementImpl.class.cast(o).value, value);
304        }
305
306        @Override
307        public int hashCode() {
308            return Objects.hash(value);
309        }
310    }
311
312    Object writeReplace() throws ObjectStreamException {
313        return new SerializableService(plugin, LocalCache.class.getName());
314    }
315}