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