16
16
import com .google .devtools .build .lib .packages .LabelPrinter ;
17
17
import com .google .devtools .build .lib .packages .Target ;
18
18
import com .google .devtools .build .lib .query2 .engine .OutputFormatterCallback ;
19
+ import com .google .devtools .build .lib .query2 .proto .proto2api .Build ;
20
+
21
+ import java .io .ByteArrayOutputStream ;
19
22
import java .io .IOException ;
20
23
import java .io .OutputStream ;
24
+ import java .util .stream .StreamSupport ;
21
25
22
26
/**
23
27
* An output formatter that outputs a protocol buffer representation of a query result and outputs
@@ -34,13 +38,70 @@ public String getName() {
34
38
public OutputFormatterCallback <Target > createPostFactoStreamCallback (
35
39
final OutputStream out , final QueryOptions options , LabelPrinter labelPrinter ) {
36
40
return new OutputFormatterCallback <Target >() {
41
+ private final LabelPrinter ourLabelPrinter = labelPrinter ;
42
+
37
43
@ Override
38
44
public void processOutput (Iterable <Target > partialResult )
39
45
throws IOException , InterruptedException {
40
- for (Target target : partialResult ) {
41
- toTargetProtoBuffer (target , labelPrinter ).writeDelimitedTo (out );
46
+ try {
47
+ StreamSupport .stream (partialResult .spliterator (), /* parallel= */ true )
48
+ .map (this ::toProto )
49
+ .map (StreamedProtoOutputFormatter ::writeDelimited )
50
+ .forEach (this ::writeToOutputStreamThreadSafe );
51
+ } catch (WrappedIOException e ) {
52
+ throw e .getCause ();
53
+ } catch (WrappedInterruptedException e ) {
54
+ throw e .getCause ();
55
+ }
56
+ }
57
+
58
+ private Build .Target toProto (Target target ) {
59
+ try {
60
+ return toTargetProtoBuffer (target , ourLabelPrinter );
61
+ } catch (InterruptedException e ) {
62
+ throw new WrappedInterruptedException (e );
63
+ }
64
+ }
65
+
66
+ private synchronized void writeToOutputStreamThreadSafe (ByteArrayOutputStream bout ) {
67
+ try {
68
+ bout .writeTo (out );
69
+ } catch (IOException e ) {
70
+ throw new RuntimeException (e );
42
71
}
43
72
}
44
73
};
45
74
}
75
+
76
+ private static ByteArrayOutputStream writeDelimited (Build .Target targetProtoBuffer ) {
77
+ try {
78
+ var bout = new ByteArrayOutputStream (targetProtoBuffer .getSerializedSize () + 10 );
79
+ targetProtoBuffer .writeDelimitedTo (bout );
80
+ return bout ;
81
+ } catch (IOException e ) {
82
+ throw new WrappedIOException (e );
83
+ }
84
+ }
85
+
86
+ private static class WrappedIOException extends RuntimeException {
87
+ private WrappedIOException (IOException cause ) {
88
+ super (cause );
89
+ }
90
+
91
+ @ Override
92
+ public synchronized IOException getCause () {
93
+ return (IOException ) super .getCause ();
94
+ }
95
+ }
96
+
97
+ private static class WrappedInterruptedException extends RuntimeException {
98
+ private WrappedInterruptedException (InterruptedException cause ) {
99
+ super (cause );
100
+ }
101
+
102
+ @ Override
103
+ public synchronized InterruptedException getCause () {
104
+ return (InterruptedException ) super .getCause ();
105
+ }
106
+ }
46
107
}
0 commit comments