001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements. See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache license, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License. You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the license for the specific language governing permissions and
015 * limitations under the license.
016 */
017package org.apache.logging.log4j.util;
018
019import java.io.ByteArrayInputStream;
020import java.io.ByteArrayOutputStream;
021import java.io.IOException;
022import java.io.InvalidObjectException;
023import java.io.ObjectInputStream;
024import java.io.ObjectOutputStream;
025import java.io.StreamCorruptedException;
026import java.lang.reflect.InvocationTargetException;
027import java.lang.reflect.Method;
028import java.lang.reflect.Modifier;
029import java.util.Arrays;
030import java.util.Collection;
031import java.util.ConcurrentModificationException;
032import java.util.HashMap;
033import java.util.Map;
034import java.util.Objects;
035
036import org.apache.logging.log4j.status.StatusLogger;
037
038/**
039 * <em>Consider this class private.</em>
040 * Array-based implementation of the {@code ReadOnlyStringMap} interface. Keys are held in a sorted array.
041 * <p>
042 * This is not a generic collection, but makes some trade-offs to optimize for the Log4j context data use case:
043 * </p>
044 * <ul>
045 *   <li>Garbage-free iteration over key-value pairs with {@code BiConsumer} and {@code TriConsumer}.</li>
046 *   <li>Fast copy. If the ThreadContextMap is also an instance of {@code SortedArrayStringMap}, the full thread context
047 *     data can be transferred with two array copies and two field updates.</li>
048 *   <li>Acceptable performance for small data sets. The current implementation stores keys in a sorted array, values
049 *     are stored in a separate array at the same index.
050 *     Worst-case performance of {@code get} and {@code containsKey} is O(log N),
051 *     worst-case performance of {@code put} and {@code remove} is O(N log N).
052 *     The expectation is that for the small values of {@code N} (less than 100) that are the vast majority of
053 *     ThreadContext use cases, the constants dominate performance more than the asymptotic performance of the
054 *     algorithms used.
055 *     </li>
056 *     <li>Compact representation.</li>
057 * </ul>
058 *
059 * @since 2.7
060 */
061public class SortedArrayStringMap implements IndexedStringMap {
062
063    /**
064     * The default initial capacity.
065     */
066    private static final int DEFAULT_INITIAL_CAPACITY = 4;
067    private static final long serialVersionUID = -5748905872274478116L;
068    private static final int HASHVAL = 31;
069
070    private static final TriConsumer<String, Object, StringMap> PUT_ALL = new TriConsumer<String, Object, StringMap>() {
071        @Override
072        public void accept(final String key, final Object value, final StringMap contextData) {
073            contextData.putValue(key, value);
074        }
075    };
076
077    /**
078     * An empty array instance to share when the table is not inflated.
079     */
080    private static final String[] EMPTY = {};
081    private static final String FROZEN = "Frozen collection cannot be modified";
082
083    private transient String[] keys = EMPTY;
084    private transient Object[] values = EMPTY;
085
086    /**
087     * The number of key-value mappings contained in this map.
088     */
089    private transient int size;
090
091    private static final Method setObjectInputFilter;
092    private static final Method getObjectInputFilter;
093    private static final Method newObjectInputFilter;
094
095    static {
096        Method[] methods = ObjectInputStream.class.getMethods();
097        Method setMethod = null;
098        Method getMethod = null;
099        for (Method method : methods) {
100            if (method.getName().equals("setObjectInputFilter")) {
101                setMethod = method;
102            } else if (method.getName().equals("getObjectInputFilter")) {
103                getMethod = method;
104            }
105        }
106        Method newMethod = null;
107        try {
108            if (setMethod != null) {
109                Class<?> clazz = Class.forName("org.apache.logging.log4j.util.internal.DefaultObjectInputFilter");
110                methods = clazz.getMethods();
111                for (Method method : methods) {
112                    if (method.getName().equals("newInstance") && Modifier.isStatic(method.getModifiers())) {
113                        newMethod = method;
114                        break;
115                    }
116                }
117            }
118        } catch (ClassNotFoundException ex) {
119            // Ignore the exception
120        }
121        newObjectInputFilter = newMethod;
122        setObjectInputFilter = setMethod;
123        getObjectInputFilter = getMethod;
124    }
125
126    /**
127     * The next size value at which to resize (capacity * load factor).
128     * @serial
129     */
130    // If table == EMPTY_TABLE then this is the initial capacity at which the
131    // table will be created when inflated.
132    private int threshold;
133    private boolean immutable;
134    private transient boolean iterating;
135
136    public SortedArrayStringMap() {
137        this(DEFAULT_INITIAL_CAPACITY);
138    }
139
140    public SortedArrayStringMap(final int initialCapacity) {
141        if (initialCapacity < 0) {
142            throw new IllegalArgumentException("Initial capacity must be at least zero but was " + initialCapacity);
143        }
144        threshold = ceilingNextPowerOfTwo(initialCapacity == 0 ? 1 : initialCapacity);
145    }
146
147    public SortedArrayStringMap(final ReadOnlyStringMap other) {
148        if (other instanceof SortedArrayStringMap) {
149            initFrom0((SortedArrayStringMap) other);
150        } else if (other != null) {
151            resize(ceilingNextPowerOfTwo(other.size()));
152            other.forEach(PUT_ALL, this);
153        }
154    }
155
156    public SortedArrayStringMap(final Map<String, ?> map) {
157        resize(ceilingNextPowerOfTwo(map.size()));
158        for (final Map.Entry<String, ?> entry : map.entrySet()) {
159            putValue(entry.getKey(), entry.getValue());
160        }
161    }
162
163    private void assertNotFrozen() {
164        if (immutable) {
165            throw new UnsupportedOperationException(FROZEN);
166        }
167    }
168
169    private void assertNoConcurrentModification() {
170        if (iterating) {
171            throw new ConcurrentModificationException();
172        }
173    }
174
175    @Override
176    public void clear() {
177        if (keys == EMPTY) {
178            return;
179        }
180        assertNotFrozen();
181        assertNoConcurrentModification();
182
183        Arrays.fill(keys, 0, size, null);
184        Arrays.fill(values, 0, size, null);
185        size = 0;
186    }
187
188    @Override
189    public boolean containsKey(final String key) {
190        return indexOfKey(key) >= 0;
191    }
192
193    @Override
194    public Map<String, String> toMap() {
195        final Map<String, String> result = new HashMap<>(size());
196        for (int i = 0; i < size(); i++) {
197            final Object value = getValueAt(i);
198            result.put(getKeyAt(i), value == null ? null : String.valueOf(value));
199        }
200        return result;
201    }
202
203    @Override
204    public void freeze() {
205        immutable = true;
206    }
207
208    @Override
209    public boolean isFrozen() {
210        return immutable;
211    }
212
213    @SuppressWarnings("unchecked")
214    @Override
215    public <V> V getValue(final String key) {
216        final int index = indexOfKey(key);
217        if (index < 0) {
218            return null;
219        }
220        return (V) values[index];
221    }
222
223    @Override
224    public boolean isEmpty() {
225        return size == 0;
226    }
227
228    @Override
229    public int indexOfKey(final String key) {
230        if (keys == EMPTY) {
231            return -1;
232        }
233        if (key == null) { // null key is located at the start of the array
234            return nullKeyIndex(); // insert at index zero
235        }
236        final int start = size > 0 && keys[0] == null ? 1 : 0;
237        return Arrays.binarySearch(keys, start, size, key);
238    }
239
240    private int nullKeyIndex() {
241        return size > 0 && keys[0] == null ? 0 : ~0;
242    }
243
244    @Override
245    public void putValue(final String key, final Object value) {
246        assertNotFrozen();
247        assertNoConcurrentModification();
248
249        if (keys == EMPTY) {
250            inflateTable(threshold);
251        }
252        final int index = indexOfKey(key);
253        if (index >= 0) {
254            keys[index] = key;
255            values[index] = value;
256        } else { // not found, so insert.
257            insertAt(~index, key, value);
258        }
259    }
260
261    private void insertAt(final int index, final String key, final Object value) {
262        ensureCapacity();
263        System.arraycopy(keys, index, keys, index + 1, size - index);
264        System.arraycopy(values, index, values, index + 1, size - index);
265        keys[index] = key;
266        values[index] = value;
267        size++;
268    }
269
270    @Override
271    public void putAll(final ReadOnlyStringMap source) {
272        if (source == this || source == null || source.isEmpty()) {
273            // this.putAll(this) does not modify this collection
274            // this.putAll(null) does not modify this collection
275            // this.putAll(empty ReadOnlyStringMap) does not modify this collection
276            return;
277        }
278        assertNotFrozen();
279        assertNoConcurrentModification();
280
281        if (source instanceof SortedArrayStringMap) {
282            if (this.size == 0) {
283                initFrom0((SortedArrayStringMap) source);
284            } else {
285                merge((SortedArrayStringMap) source);
286            }
287        } else if (source != null) {
288            source.forEach(PUT_ALL, this);
289        }
290    }
291
292    private void initFrom0(final SortedArrayStringMap other) {
293        if (keys.length < other.size) {
294            keys = new String[other.threshold];
295            values = new Object[other.threshold];
296        }
297        System.arraycopy(other.keys, 0, keys, 0, other.size);
298        System.arraycopy(other.values, 0, values, 0, other.size);
299
300        size = other.size;
301        threshold = other.threshold;
302    }
303
304    private void merge(final SortedArrayStringMap other) {
305        final String[] myKeys = keys;
306        final Object[] myVals = values;
307        final int newSize = other.size + this.size;
308        threshold = ceilingNextPowerOfTwo(newSize);
309        if (keys.length < threshold) {
310            keys = new String[threshold];
311            values = new Object[threshold];
312        }
313        // move largest collection to the beginning of this data structure, smallest to the end
314        boolean overwrite = true;
315        if (other.size() > size()) {
316            // move original key-values to end
317            System.arraycopy(myKeys, 0, keys, other.size, this.size);
318            System.arraycopy(myVals, 0, values, other.size, this.size);
319
320            // insert additional key-value pairs at the beginning
321            System.arraycopy(other.keys, 0, keys, 0, other.size);
322            System.arraycopy(other.values, 0, values, 0, other.size);
323            size = other.size;
324
325            // loop over original keys and insert (drop values for same key)
326            overwrite = false;
327        } else {
328            System.arraycopy(myKeys, 0, keys, 0, this.size);
329            System.arraycopy(myVals, 0, values, 0, this.size);
330
331            // move additional key-value pairs to end
332            System.arraycopy(other.keys, 0, keys, this.size, other.size);
333            System.arraycopy(other.values, 0, values, this.size, other.size);
334
335            // new values are at the end, will be processed below. Overwrite is true.
336        }
337        for (int i = size; i < newSize; i++) {
338            final int index = indexOfKey(keys[i]);
339            if (index < 0) { // key does not exist
340                insertAt(~index, keys[i], values[i]);
341            } else if (overwrite) { // existing key: only overwrite when looping over the new values
342                keys[index] = keys[i];
343                values[index] = values[i];
344            }
345        }
346        // prevent memory leak: null out references
347        Arrays.fill(keys, size, newSize, null);
348        Arrays.fill(values, size, newSize, null);
349    }
350
351    private void ensureCapacity() {
352        if (size >= threshold) {
353            resize(threshold * 2);
354        }
355    }
356
357    private void resize(final int newCapacity) {
358        final String[] oldKeys = keys;
359        final Object[] oldValues = values;
360
361        keys = new String[newCapacity];
362        values = new Object[newCapacity];
363
364        System.arraycopy(oldKeys, 0, keys, 0, size);
365        System.arraycopy(oldValues, 0, values, 0, size);
366
367        threshold = newCapacity;
368    }
369
370    /**
371     * Inflates the table.
372     */
373    private void inflateTable(final int toSize) {
374        threshold = toSize;
375        keys = new String[toSize];
376        values = new Object[toSize];
377    }
378
379    @Override
380    public void remove(final String key) {
381        if (keys == EMPTY) {
382            return;
383        }
384
385        final int index = indexOfKey(key);
386        if (index >= 0) {
387            assertNotFrozen();
388            assertNoConcurrentModification();
389
390            System.arraycopy(keys, index + 1, keys, index, size - 1 - index);
391            System.arraycopy(values, index + 1, values, index, size - 1 - index);
392            keys[size - 1] = null;
393            values[size - 1] = null;
394            size--;
395        }
396    }
397
398    @Override
399    public String getKeyAt(final int index) {
400        if (index < 0 || index >= size) {
401            return null;
402        }
403        return keys[index];
404    }
405
406    @SuppressWarnings("unchecked")
407    @Override
408    public <V> V getValueAt(final int index) {
409        if (index < 0 || index >= size) {
410            return null;
411        }
412        return (V) values[index];
413    }
414
415    @Override
416    public int size() {
417        return size;
418    }
419
420    @SuppressWarnings("unchecked")
421    @Override
422    public <V> void forEach(final BiConsumer<String, ? super V> action) {
423        iterating = true;
424        try {
425            for (int i = 0; i < size; i++) {
426                action.accept(keys[i], (V) values[i]);
427            }
428        } finally {
429            iterating = false;
430        }
431    }
432
433    @SuppressWarnings("unchecked")
434    @Override
435    public <V, T> void forEach(final TriConsumer<String, ? super V, T> action, final T state) {
436        iterating = true;
437        try {
438            for (int i = 0; i < size; i++) {
439                action.accept(keys[i], (V) values[i], state);
440            }
441        } finally {
442            iterating = false;
443        }
444    }
445
446    @Override
447    public boolean equals(final Object obj) {
448        if (obj == this) {
449            return true;
450        }
451        if (!(obj instanceof SortedArrayStringMap)) {
452            return false;
453        }
454        final SortedArrayStringMap other = (SortedArrayStringMap) obj;
455        if (this.size() != other.size()) {
456            return false;
457        }
458        for (int i = 0; i < size(); i++) {
459            if (!Objects.equals(keys[i], other.keys[i])) {
460                return false;
461            }
462            if (!Objects.equals(values[i], other.values[i])) {
463                return false;
464            }
465        }
466        return true;
467    }
468
469    @Override
470    public int hashCode() {
471        int result = 37;
472        result = HASHVAL * result + size;
473        result = HASHVAL * result + hashCode(keys, size);
474        result = HASHVAL * result + hashCode(values, size);
475        return result;
476    }
477
478    private static int hashCode(final Object[] values, final int length) {
479        int result = 1;
480        for (int i = 0; i < length; i++) {
481            result = HASHVAL * result + (values[i] == null ? 0 : values[i].hashCode());
482        }
483        return result;
484    }
485
486    @Override
487    public String toString() {
488        final StringBuilder sb = new StringBuilder(256);
489        sb.append('{');
490        for (int i = 0; i < size; i++) {
491            if (i > 0) {
492                sb.append(", ");
493            }
494            sb.append(keys[i]).append('=');
495            sb.append(values[i] == this ? "(this map)" : values[i]);
496        }
497        sb.append('}');
498        return sb.toString();
499    }
500
501    /**
502     * Save the state of the {@code SortedArrayStringMap} instance to a stream (i.e.,
503     * serialize it).
504     *
505     * @serialData The <i>capacity</i> of the SortedArrayStringMap (the length of the
506     *             bucket array) is emitted (int), followed by the
507     *             <i>size</i> (an int, the number of key-value
508     *             mappings), followed by the key (Object) and value (Object)
509     *             for each key-value mapping.  The key-value mappings are
510     *             emitted in no particular order.
511     */
512    private void writeObject(final java.io.ObjectOutputStream s) throws IOException {
513        // Write out the threshold, and any hidden stuff
514        s.defaultWriteObject();
515
516        // Write out number of buckets
517        if (keys == EMPTY) {
518            s.writeInt(ceilingNextPowerOfTwo(threshold));
519        } else {
520            s.writeInt(keys.length);
521        }
522
523        // Write out size (number of Mappings)
524        s.writeInt(size);
525
526        // Write out keys and values (alternating)
527        if (size > 0) {
528            for (int i = 0; i < size; i++) {
529                s.writeObject(keys[i]);
530                try {
531                    s.writeObject(marshall(values[i]));
532                } catch (final Exception e) {
533                    handleSerializationException(e, i, keys[i]);
534                    s.writeObject(null);
535                }
536            }
537        }
538    }
539
540    private static byte[] marshall(final Object obj) throws IOException {
541        if (obj == null) {
542            return null;
543        }
544        final ByteArrayOutputStream bout = new ByteArrayOutputStream();
545        try (ObjectOutputStream oos = new ObjectOutputStream(bout)) {
546            oos.writeObject(obj);
547            oos.flush();
548            return bout.toByteArray();
549        }
550    }
551
552    private static Object unmarshall(final byte[] data, ObjectInputStream inputStream)
553            throws IOException, ClassNotFoundException {
554        final ByteArrayInputStream bin = new ByteArrayInputStream(data);
555        Collection<String> allowedClasses = null;
556        ObjectInputStream ois;
557        if (inputStream instanceof FilteredObjectInputStream) {
558            allowedClasses = ((FilteredObjectInputStream) inputStream).getAllowedClasses();
559            ois = new FilteredObjectInputStream(bin, allowedClasses);
560        } else {
561            try {
562                Object obj = getObjectInputFilter.invoke(inputStream);
563                Object filter = newObjectInputFilter.invoke(null, obj);
564                ois = new ObjectInputStream(bin);
565                setObjectInputFilter.invoke(ois, filter);
566            } catch (IllegalAccessException | InvocationTargetException ex) {
567                throw new StreamCorruptedException("Unable to set ObjectInputFilter on stream");
568            }
569        }
570        try {
571            return ois.readObject();
572        } finally {
573            ois.close();
574        }
575    }
576
577    /**
578     * Calculate the next power of 2, greater than or equal to x.
579     * <p>
580     * From Hacker's Delight, Chapter 3, Harry S. Warren Jr.
581     *
582     * @param x Value to round up
583     * @return The next power of 2 from x inclusive
584     */
585    private static int ceilingNextPowerOfTwo(final int x) {
586        final int BITS_PER_INT = 32;
587        return 1 << (BITS_PER_INT - Integer.numberOfLeadingZeros(x - 1));
588    }
589
590    /**
591     * Reconstitute the {@code SortedArrayStringMap} instance from a stream (i.e.,
592     * deserialize it).
593     */
594    private void readObject(final java.io.ObjectInputStream s)  throws IOException, ClassNotFoundException {
595        if (!(s instanceof FilteredObjectInputStream) && setObjectInputFilter == null) {
596            throw new IllegalArgumentException("readObject requires a FilteredObjectInputStream or an ObjectInputStream that accepts an ObjectInputFilter");
597        }
598        // Read in the threshold (ignored), and any hidden stuff
599        s.defaultReadObject();
600
601        // set other fields that need values
602        keys = EMPTY;
603        values = EMPTY;
604
605        // Read in number of buckets
606        final int capacity = s.readInt();
607        if (capacity < 0) {
608            throw new InvalidObjectException("Illegal capacity: " + capacity);
609        }
610
611        // Read number of mappings
612        final int mappings = s.readInt();
613        if (mappings < 0) {
614            throw new InvalidObjectException("Illegal mappings count: " + mappings);
615        }
616
617        // allocate the bucket array;
618        if (mappings > 0) {
619            inflateTable(capacity);
620        } else {
621            threshold = capacity;
622        }
623
624        // Read the keys and values, and put the mappings in the arrays
625        for (int i = 0; i < mappings; i++) {
626            keys[i] = (String) s.readObject();
627            try {
628                final byte[] marshalledObject = (byte[]) s.readObject();
629                values[i] = marshalledObject == null ? null : unmarshall(marshalledObject, s);
630            } catch (final Exception | LinkageError error) {
631                handleSerializationException(error, i, keys[i]);
632                values[i] = null;
633            }
634        }
635        size = mappings;
636    }
637
638    private void handleSerializationException(final Throwable t, final int i, final String key) {
639        StatusLogger.getLogger().warn("Ignoring {} for key[{}] ('{}')", String.valueOf(t), i, keys[i]);
640    }
641}