Солнышко Ксения, ИТМО ФИТиП, Cluster range requests

src/main/java/ru/vk/itmo/test/solnyshkoksenia/CustomHttpSession.java

@@ -0,0 +1,96 @@
+package ru.vk.itmo.test.solnyshkoksenia;
+
+import one.nio.http.HttpServer;
+import one.nio.http.HttpSession;
+import one.nio.http.Response;
+import one.nio.net.Socket;
+import ru.vk.itmo.dao.Entry;
+
+import java.io.IOException;
+import java.lang.foreign.MemorySegment;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Iterator;
+
+public class CustomHttpSession extends HttpSession {
+    private static final byte[] CRLF = "\r\n".getBytes(StandardCharsets.UTF_8);
+    private static final byte[] LF = "\n".getBytes(StandardCharsets.UTF_8);
+    private static final byte[] EMPTY_CHUNK = "0\r\n\r\n".getBytes(StandardCharsets.UTF_8);
+    Iterator<Entry<MemorySegment>> iterator;
+
+    public CustomHttpSession(Socket socket, HttpServer server) {
+        super(socket, server);
+    }
+
+    @Override
+    protected void processWrite() throws Exception {
+        super.processWrite();
+        nextChunk();
+    }
+
+    public void stream(Iterator<Entry<MemorySegment>> iterator) throws IOException {
+        this.iterator = iterator;
+        Response response = new Response(Response.OK);
+        response.addHeader("Transfer-Encoding: chunked");
+        writeResponse(response, false);
+
+        nextChunk();
+    }
+
+    private void nextChunk() throws IOException {
+        while (iterator.hasNext() && queueHead == null) {
+            Entry<MemorySegment> next = iterator.next();
+            ByteBuffer key = next.key().asByteBuffer();
+            ByteBuffer value = next.value().asByteBuffer();
+            int payloadSize = key.remaining() + value.remaining() + LF.length;
+            String payloadSizeStr = Integer.toHexString(payloadSize);
+            byte[] payloadSizeStrBytes = payloadSizeStr.getBytes(StandardCharsets.UTF_8);
+            write(payloadSizeStrBytes, 0, payloadSizeStrBytes.length);
+            write(CRLF, 0, CRLF.length);
+            write(new ReferenceQueueItem(key));
+            write(LF, 0, LF.length);
+            write(new ReferenceQueueItem(value));
+            write(CRLF, 0, CRLF.length);
+        }
+
+        if (!iterator.hasNext()) {
+            write(EMPTY_CHUNK, 0, EMPTY_CHUNK.length);
+
+            if ((this.handling = pipeline.pollFirst()) != null) {
+                if (handling == FIN) {
+                    scheduleClose();
+                } else {
+                    server.handleRequest(handling, this);
+                }
+            }
+        }
+    }
+
+    public void sendError(Throwable e) {
+        log.error("Exception during handleRequest", e);
+        try {
+            sendResponse(new Response(Response.INTERNAL_ERROR, Response.EMPTY));
+        } catch (IOException ex) {
+            log.error("Exception while sending close connection", e);
+            scheduleClose();
+        }
+    }
+
+    static class ReferenceQueueItem extends QueueItem {
+        private final ByteBuffer buffer;
+
+        ReferenceQueueItem(ByteBuffer buffer) {
+            this.buffer = buffer;
+        }
+
+        @Override
+        public int remaining() {
+            return buffer.remaining();
+        }
+
+        @Override
+        public int write(Socket socket) throws IOException {
+            return socket.write(buffer);
+        }
+    }
+}

src/main/java/ru/vk/itmo/test/solnyshkoksenia/CustomSubscriber.java

@@ -0,0 +1,126 @@
+package ru.vk.itmo.test.solnyshkoksenia;
+
+import java.net.http.HttpResponse.BodySubscriber;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Flow;
+
+public class CustomSubscriber implements BodySubscriber<byte[]> {
+    private static final byte DELIMITER = '\n';
+    volatile CompletableFuture<byte[]> bodyCF;
+    Flow.Subscription subscription;
+    List<ByteBuffer> responseData = new CopyOnWriteArrayList<>();
+
+    @Override
+    public CompletionStage<byte[]> getBody() {
+        while (bodyCF == null) {
+            Thread.onSpinWait();
+        }
+        return bodyCF;
+    }
+
+    @Override
+    public void onSubscribe(Flow.Subscription subscription) {
+        this.subscription = subscription;
+        subscription.request(1);
+    }
+
+    @Override
+    public void onNext(List<ByteBuffer> buffers) {
+        buffers.forEach(ByteBuffer::rewind);
+        responseData.addAll(buffers);
+        subscription.request(1);
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+        bodyCF = CompletableFuture.failedFuture(throwable);
+    }
+
+    @Override
+    public void onComplete() {
+        bodyCF = CompletableFuture.completedFuture(toBytes(responseData));
+    }
+
+    private List<byte[]> toArrays(List<ByteBuffer> buffers) {
+        List<byte[]> chunks = new ArrayList<>();
+        for (ByteBuffer buffer : buffers) {
+            int remaining = buffer.remaining();
+            byte[] cur = new byte[remaining];
+            buffer.get(cur, 0, remaining);
+            chunks.add(cur);
+        }
+        return chunks;
+    }
+
+    private boolean containsDelim(byte[] bytes) {
+        for (int element : bytes) {
+            if (element == DELIMITER) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private boolean startsWithDelim(byte[] bytes) {
+        return DELIMITER == bytes[0];
+    }
+
+    private byte[] merge(byte[] src1, byte[] src2) {
+        byte[] dst = new byte[src1.length + src2.length];
+        System.arraycopy(src1, 0, dst, 0, src1.length);
+        System.arraycopy(src2, 0, dst, src1.length, src2.length);
+        return dst;
+    }
+
+    private byte[] toArray(List<byte[]> bytes) {
+        var size = bytes.stream().mapToInt(b -> b.length).sum() + bytes.size();
+        byte[] dst = new byte[size];
+        int offset = 0;
+        for (byte[] src : bytes) {
+            System.arraycopy(src, 0, dst, offset, src.length);
+            offset += src.length;
+            dst[offset] = DELIMITER;
+            offset++;
+        }
+        return dst;
+    }
+
+    private byte[] toBytes(List<ByteBuffer> buffers) {
+        List<byte[]> chunks = toArrays(buffers);
+        List<byte[]> bytes = new ArrayList<>();
+
+        boolean predEndsWithDelim = false;
+        boolean predContainsDelim = true;
+
+        byte[] pred = null;
+        for (int i = 0; i < chunks.size(); i++) {
+            byte[] cur = chunks.get(i);
+            byte[] next = i + 1 == chunks.size() ? new byte[0] : chunks.get(i + 1);
+            if (startsWithDelim(cur) && !predContainsDelim) {
+                cur = merge(pred, cur);
+            } else if (!containsDelim(cur) && predEndsWithDelim) {
+                if (startsWithDelim(next)) {
+                    bytes.add(pred);
+                } else {
+                    cur = merge(pred, cur);
+                }
+            } else if (pred != null) {
+                bytes.add(pred);
+            }
+
+            predEndsWithDelim = DELIMITER == cur[cur.length - 1];
+            predContainsDelim = containsDelim(cur);
+            pred = cur;
+        }
+
+        if (pred != null) {
+            bytes.add(pred);
+        }
+        return toArray(bytes);
+    }
+}

src/main/java/ru/vk/itmo/test/solnyshkoksenia/MergeRangeResult.java

@@ -0,0 +1,75 @@
+package ru.vk.itmo.test.solnyshkoksenia;
+
+import one.nio.http.Response;
+import ru.vk.itmo.dao.BaseEntry;
+import ru.vk.itmo.dao.Entry;
+import ru.vk.itmo.test.solnyshkoksenia.dao.MemorySegmentComparator;
+import ru.vk.itmo.test.solnyshkoksenia.dao.MergeIterator;
+
+import java.lang.foreign.MemorySegment;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+public final class MergeRangeResult {
+    private static final Comparator<MemorySegment> comparator = new MemorySegmentComparator();
+
+    private MergeRangeResult() {
+    }
+
+    public static Iterator<Entry<MemorySegment>> range(Iterator<Entry<MemorySegment>> firstIterator,
+                                                       List<Response> responses) {
+        List<Iterator<Entry<MemorySegment>>> iterators = new ArrayList<>(responses.size() + 1);
+        for (var response : responses) {
+            iterators.add(iterator(response));
+        }
+        iterators.add(firstIterator);
+
+        return new MergeIterator<>(iterators, (e1, e2) -> comparator.compare(e1.key(), e2.key())) {
+            @Override
+            protected boolean skip(Entry<MemorySegment> memorySegmentEntry) {
+                return memorySegmentEntry.value() == null;
+            }
+        };
+    }
+
+    private static Iterator<Entry<MemorySegment>> iterator(Response response) {
+        byte[] body = response.getBody();
+        char separator = '\n';
+        return new Iterator<>() {
+            int offset;
+
+            @Override
+            public boolean hasNext() {
+                return offset < body.length;
+            }
+
+            @Override
+            public Entry<MemorySegment> next() {
+                if (!hasNext()) {
+                    throw new NoSuchElementException();
+                }
+                MemorySegment key = getMS();
+                MemorySegment value = getMS();
+                return new BaseEntry<>(key, value);
+            }
+
+            private MemorySegment getMS() {
+                int startOffset = offset;
+
+                byte b = body[offset];
+                while (b != separator) {
+                    b = body[++offset];
+                }
+
+                byte[] tmp = new byte[offset - startOffset];
+                System.arraycopy(body, startOffset, tmp, 0, offset - startOffset);
+
+                offset++;
+                return MemorySegment.ofArray(tmp);
+            }
+        };
+    }
+}