Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package com.facebook.presto.operator;
 
 
 import java.net.URI;
 import java.util.List;
 
 import static com.facebook.presto.PrestoMediaTypes.PRESTO_PAGES;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_PAGE_NEXT_TOKEN;
 import static com.facebook.presto.client.PrestoHeaders.PRESTO_PAGE_TOKEN;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
         implements Function<RequestResponse>
 {
     private final LoadingCache<URIMockBufferbuffers = CacheBuilder.newBuilder().build(new CacheLoader<URIMockBuffer>()
     {
         @Override
         public MockBuffer load(URI location)
         {
             return new MockBuffer(location);
         }
     });
 
     private final DataSize expectedMaxSize;
 
     public MockExchangeRequestProcessor(DataSize expectedMaxSize)
     {
         this. = expectedMaxSize;
     }
 
     public void addPage(URI locationPage page)
     {
         .getUnchecked(location).addPage(page);
     }
 
     public void setComplete(URI location)
     {
         .getUnchecked(location).setCompleted();
     }
 
     @Override
     public Response apply(Request request)
     {
         if (request.getMethod().equalsIgnoreCase("DELETE")) {
             return new TestingResponse(., ImmutableListMultimap.<StringString>of(), new byte[0]);
         }
 
         // verify we got a data size and it parses correctly
         DataSize maxSize = DataSize.valueOf(request.getHeader(.));
         assertEquals(maxSize);
 
         RequestLocation requestLocation = new RequestLocation(request.getUri());
         URI location = requestLocation.getLocation();
 
         BufferResult result = .getUnchecked(location).getPages(requestLocation.getSequenceId(), maxSize);
         List<Pagepages = result.getPages();
 
         byte[] bytes = new byte[0];
         HttpStatus status;
        if (!pages.isEmpty()) {
            DynamicSliceOutput sliceOutput = new DynamicSliceOutput(64);
            PagesSerde.writePages(new BlockEncodingManager(new TypeRegistry()), sliceOutputpages);
            bytes = sliceOutput.slice().getBytes();
            status = .;
        }
        else if (result.isBufferClosed()) {
            status = .;
        }
        else {
            status = .;
        }
        return new TestingResponse(
                status,
                ImmutableListMultimap.of(
                        ,
                        , String.valueOf(result.getToken()),
                        , String.valueOf(result.getNextToken())
                ),
                bytes);
    }
    private class RequestLocation
    {
        private final URI location;
        private final long sequenceId;
        public RequestLocation(URI uri)
        {
            String string = uri.toString();
            int index = string.lastIndexOf('/');
             = URI.create(string.substring(0, index));
             = Long.parseLong(string.substring(index + 1));
        }
        public URI getLocation()
        {
            return ;
        }
        public long getSequenceId()
        {
            return ;
        }
    }
    private static class MockBuffer
    {
        private final URI location;
        private final AtomicBoolean completed = new AtomicBoolean();
        private final AtomicLong token = new AtomicLong();
        private final BlockingQueue<Pagepages = new LinkedBlockingQueue<>();
        private MockBuffer(URI location)
        {
            this. = location;
        }
        public void setCompleted()
        {
            .set(true);
        }
        public synchronized void addPage(Page page)
        {
            checkState(.get() != ."Location %s is complete");
            .add(page);
        }
        public BufferResult getPages(long sequenceIdDataSize maxSize)
        {
            // if location is complete return GONE
            if (.get() && .isEmpty()) {
                return BufferResult.emptyResults(.get(), true);
            }
            assertEquals(sequenceId.get(), "token");
            // wait for a single page to arrive
            Page page = null;
            try {
                page = .poll(10, .);
            }
            catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
            // if no page, return NO CONTENT
            if (page == null) {
                return BufferResult.emptyResults(.get(), false);
            }
            // add pages up to the size limit
            List<PageresponsePages = new ArrayList<>();
            responsePages.add(page);
            long responseSize = page.getSizeInBytes();
            while (responseSize < maxSize.toBytes()) {
                page = .poll();
                if (page == null) {
                    break;
                }
                responsePages.add(page);
                responseSize += page.getSizeInBytes();
            }
            // update sequence id
            long nextToken = .get() + responsePages.size();
            BufferResult bufferResult = new BufferResult(.get(), nextTokenfalseresponsePages);
            .set(nextToken);
            return bufferResult;
        }
    }
New to GrepCode? Check out our FAQ X