1+ package graphql.execution.instrumentation.dataloader
2+
3+ import graphql.ExecutionInput
4+ import graphql.ExecutionResult
5+ import graphql.GraphQL
6+ import graphql.TestUtil
7+ import graphql.execution.Async
8+ import graphql.schema.DataFetcher
9+ import graphql.schema.DataFetchingEnvironment
10+ import graphql.schema.idl.RuntimeWiring
11+ import org.apache.commons.lang3.concurrent.BasicThreadFactory
12+ import org.dataloader.BatchLoader
13+ import org.dataloader.DataLoader
14+ import org.dataloader.DataLoaderOptions
15+ import org.dataloader.DataLoaderRegistry
16+ import spock.lang.Specification
17+
18+ import java.util.concurrent.CompletableFuture
19+ import java.util.concurrent.CompletionStage
20+ import java.util.concurrent.SynchronousQueue
21+ import java.util.concurrent.ThreadFactory
22+ import java.util.concurrent.ThreadPoolExecutor
23+ import java.util.concurrent.TimeUnit
24+
25+ import static graphql.schema.idl.TypeRuntimeWiring.newTypeWiring
26+
27+ class DataLoaderHangingTest extends Specification {
28+
29+ public static final int NUM_OF_REPS = 50
30+
31+ def " deadlock attempt"
32+ setup :
33+ def sdl = """
34+ type Album {
35+ id: ID!
36+ title: String!
37+ artist: Artist
38+ songs(
39+ limit: Int,
40+ nextToken: String
41+ ): ModelSongConnection
42+ }
43+
44+ type Artist {
45+ id: ID!
46+ name: String!
47+ albums(
48+ limit: Int,
49+ nextToken: String
50+ ): ModelAlbumConnection
51+ songs(
52+ limit: Int,
53+ nextToken: String
54+ ): ModelSongConnection
55+ }
56+
57+ type ModelAlbumConnection {
58+ items: [Album]
59+ nextToken: String
60+ }
61+
62+ type ModelArtistConnection {
63+ items: [Artist]
64+ nextToken: String
65+ }
66+
67+ type ModelSongConnection {
68+ items: [Song]
69+ nextToken: String
70+ }
71+
72+ type Query {
73+ listArtists(limit: Int, nextToken: String): ModelArtistConnection
74+ }
75+
76+ type Song {
77+ id: ID!
78+ title: String!
79+ artist: Artist
80+ album: Album
81+ }
82+ """
83+
84+ ThreadFactory threadFactory = new BasicThreadFactory.Builder ()
85+ .namingPattern(" resolver-chain-thread-%d" . build()
86+ def executor = new ThreadPoolExecutor (15 , 15 , 0L ,
87+ TimeUnit . MILLISECONDS , new SynchronousQueue<> (), threadFactory,
88+ new ThreadPoolExecutor.CallerRunsPolicy ())
89+
90+ def dataLoaderAlbums = new DataLoader<Object , Object > (new BatchLoader<DataFetchingEnvironment , List<Object > > () {
91+ @Override
92+ CompletionStage<List<List<Object > > > load (List<DataFetchingEnvironment > keys ) {
93+ return CompletableFuture . supplyAsync({
94+ def limit = keys. first(). getArgument(" limit" as Integer
95+ return keys. collect({ k ->
96+ def albums = []
97+ for (int i = 1 ; i <= limit; i++ ) {
98+ albums. add([' id' " artist-$k. source . id -$i " ' title' " album-$i "
99+ }
100+ def albumsConnection = [' nextToken' ' album-next' ' items'
101+ return albumsConnection
102+ })
103+ }, executor)
104+ }
105+ }, DataLoaderOptions . newOptions(). setMaxBatchSize(5 ))
106+
107+ def dataLoaderSongs = new DataLoader<Object , Object > (new BatchLoader<DataFetchingEnvironment , List<Object > > () {
108+ @Override
109+ CompletionStage<List<List<Object > > > load (List<DataFetchingEnvironment > keys ) {
110+ return CompletableFuture . supplyAsync({
111+ def limit = keys. first(). getArgument(" limit" as Integer
112+ return keys. collect({ k ->
113+ def songs = []
114+ for (int i = 1 ; i <= limit; i++ ) {
115+ songs. add([' id' " album-$k. source . id -$i " ' title' " song-$i "
116+ }
117+ def songsConnection = [' nextToken' ' song-next' ' items'
118+ return songsConnection
119+ })
120+ }, executor)
121+ }
122+ }, DataLoaderOptions . newOptions(). setMaxBatchSize(5 ))
123+
124+ def dataLoaderRegistry = new DataLoaderRegistry ()
125+ dataLoaderRegistry. register(" artist.albums"
126+ dataLoaderRegistry. register(" album.songs"
127+
128+
129+ def albumsDf = new MyForwardingDataFetcher (dataLoaderAlbums)
130+ def songsDf = new MyForwardingDataFetcher (dataLoaderSongs)
131+
132+ def dataFetcherArtists = new DataFetcher () {
133+ @Override
134+ Object get (DataFetchingEnvironment environment ) {
135+ def limit = environment. getArgument(" limit" as Integer
136+ def artists = []
137+ for (int i = 1 ; i <= limit; i++ ) {
138+ artists. add([' id' " artist-$i " ' name' " artist-$i "
139+ }
140+ return [' nextToken' ' artist-next' ' items'
141+ }
142+ }
143+
144+ def wiring = RuntimeWiring . newRuntimeWiring()
145+ .type(newTypeWiring(" Query"
146+ .dataFetcher(" listArtists"
147+ .type(newTypeWiring(" Artist"
148+ .dataFetcher(" albums"
149+ .type(newTypeWiring(" Album"
150+ .dataFetcher(" songs"
151+ .build()
152+
153+ def schema = TestUtil . schema(sdl, wiring)
154+
155+ when :
156+ def graphql = GraphQL . newGraphQL(schema)
157+ .instrumentation(new DataLoaderDispatcherInstrumentation (dataLoaderRegistry))
158+ .build()
159+
160+ then : " execution shouldn't hang"
161+ List<CompletableFuture<ExecutionResult > > futures = []
162+ for (int i = 0 ; i < NUM_OF_REPS ; i++ ) {
163+ def result = graphql. executeAsync(ExecutionInput . newExecutionInput()
164+ .query("""
165+ query getArtistsWithData {
166+ listArtists(limit: 1) {
167+ items {
168+ name
169+ albums(limit: 200) {
170+ items {
171+ title
172+ # Uncommenting the following causes query to timeout
173+ songs(limit: 5) {
174+ nextToken
175+ items {
176+ title
177+ }
178+ }
179+ }
180+ }
181+ }
182+ }
183+ }
184+ """ )
185+ .build())
186+ result. whenComplete({ res , error ->
187+ if (error) {
188+ throw error
189+ }
190+ assert res. errors. empty
191+ })
192+ // add all futures
193+ futures. add(result)
194+ }
195+ // wait for each future to complete and grab the results
196+ Async . each(futures)
197+ .whenComplete({ results , error ->
198+ if (error) {
199+ throw error
200+ }
201+ results. each { assert it. errors. empty }
202+ })
203+ .join()
204+ }
205+
206+ static class MyForwardingDataFetcher implements DataFetcher<CompletableFuture<Object > > {
207+
208+ private final DataLoader dataLoader
209+
210+ public MyForwardingDataFetcher (DataLoader dataLoader) {
211+ this . dataLoader = dataLoader
212+ }
213+
214+ @Override
215+ CompletableFuture<Object > get(DataFetchingEnvironment environment) {
216+ return dataLoader. load(environment)
217+ }
218+ }
219+ }
0 commit comments