|
1 | 1 | package tech.ydb.topic.write.impl; |
2 | 2 |
|
| 3 | +import java.io.IOException; |
| 4 | +import java.time.Instant; |
| 5 | +import java.util.List; |
3 | 6 | import java.util.concurrent.CompletableFuture; |
4 | | -import java.util.concurrent.atomic.AtomicBoolean; |
| 7 | +import java.util.stream.Collectors; |
| 8 | + |
| 9 | +import com.google.protobuf.UnsafeByteOperations; |
| 10 | +import org.slf4j.Logger; |
| 11 | +import org.slf4j.LoggerFactory; |
5 | 12 |
|
6 | 13 | import tech.ydb.common.transaction.YdbTransaction; |
| 14 | +import tech.ydb.core.utils.ProtobufUtils; |
| 15 | +import tech.ydb.proto.topic.YdbTopic; |
| 16 | +import tech.ydb.topic.description.Codec; |
| 17 | +import tech.ydb.topic.description.MetadataItem; |
7 | 18 | import tech.ydb.topic.settings.SendSettings; |
| 19 | +import tech.ydb.topic.utils.Encoder; |
8 | 20 | import tech.ydb.topic.write.Message; |
9 | 21 | import tech.ydb.topic.write.WriteAck; |
10 | 22 |
|
11 | 23 | public class EnqueuedMessage { |
12 | | - private final Message message; |
| 24 | + |
| 25 | + // use logger from WriterImpl |
| 26 | + private static final Logger logger = LoggerFactory.getLogger(WriterImpl.class); |
| 27 | + |
| 28 | + private Long seqNo; |
| 29 | + private byte[] bytes; |
| 30 | + private final long originLength; |
| 31 | + private final Instant createdAt; |
| 32 | + private final List<MetadataItem> items; |
| 33 | + |
13 | 34 | private final CompletableFuture<WriteAck> future = new CompletableFuture<>(); |
14 | | - private final AtomicBoolean isCompressed = new AtomicBoolean(); |
15 | | - private final AtomicBoolean isProcessingFailed = new AtomicBoolean(); |
16 | | - private final long uncompressedSizeBytes; |
17 | 35 | private final YdbTransaction transaction; |
18 | | - private long compressedSizeBytes; |
19 | | - private Long seqNo; |
20 | 36 |
|
21 | | - public EnqueuedMessage(Message message, SendSettings sendSettings) { |
22 | | - this.message = message; |
23 | | - this.uncompressedSizeBytes = message.getData().length; |
24 | | - this.transaction = sendSettings != null ? sendSettings.getTransaction() : null; |
25 | | - } |
| 37 | + private volatile boolean isReady = false; |
| 38 | + private volatile IOException compressError = null; |
26 | 39 |
|
27 | | - public Message getMessage() { |
28 | | - return message; |
29 | | - } |
| 40 | + public EnqueuedMessage(Message message, SendSettings sendSettings, boolean noCompression) { |
| 41 | + this.bytes = message.getData(); |
| 42 | + this.createdAt = message.getCreateTimestamp(); |
| 43 | + this.items = message.getMetadataItems(); |
| 44 | + this.seqNo = message.getSeqNo(); |
30 | 45 |
|
31 | | - public CompletableFuture<WriteAck> getFuture() { |
32 | | - return future; |
33 | | - } |
34 | | - |
35 | | - public boolean isCompressed() { |
36 | | - return isCompressed.get(); |
| 46 | + this.originLength = bytes.length; |
| 47 | + this.transaction = sendSettings != null ? sendSettings.getTransaction() : null; |
| 48 | + this.isReady = noCompression; |
37 | 49 | } |
38 | 50 |
|
39 | | - public void setCompressed(boolean compressed) { |
40 | | - this.isCompressed.set(compressed); |
| 51 | + public boolean isReady() { |
| 52 | + return isReady; |
41 | 53 | } |
42 | 54 |
|
43 | | - public boolean isProcessingFailed() { |
44 | | - return isProcessingFailed.get(); |
| 55 | + public long getOriginalSize() { |
| 56 | + return originLength; |
45 | 57 | } |
46 | 58 |
|
47 | | - public void setProcessingFailed(boolean processingFailed) { |
48 | | - isProcessingFailed.set(processingFailed); |
49 | | - } |
50 | | - |
51 | | - public long getUncompressedSizeBytes() { |
52 | | - return uncompressedSizeBytes; |
| 59 | + public long getSize() { |
| 60 | + return bytes.length; |
53 | 61 | } |
54 | 62 |
|
55 | | - public long getCompressedSizeBytes() { |
56 | | - return compressedSizeBytes; |
| 63 | + public IOException getCompressError() { |
| 64 | + return compressError; |
57 | 65 | } |
58 | 66 |
|
59 | | - public void setCompressedSizeBytes(long compressedSizeBytes) { |
60 | | - this.compressedSizeBytes = compressedSizeBytes; |
| 67 | + public void encode(String writeId, Codec codec) { |
| 68 | + logger.trace("[{}] Started encoding message", writeId); |
| 69 | + |
| 70 | + try { |
| 71 | + bytes = Encoder.encode(codec, bytes); |
| 72 | + isReady = true; |
| 73 | + logger.trace("[{}] Successfully finished encoding message", writeId); |
| 74 | + } catch (IOException ex) { |
| 75 | + logger.error("[{}] Exception while encoding message: ", writeId, ex); |
| 76 | + isReady = true; |
| 77 | + future.completeExceptionally(ex); |
| 78 | + } |
61 | 79 | } |
62 | 80 |
|
63 | | - public long getSizeBytes() { |
64 | | - return isCompressed() ? getCompressedSizeBytes() : getUncompressedSizeBytes(); |
| 81 | + public CompletableFuture<WriteAck> getFuture() { |
| 82 | + return future; |
65 | 83 | } |
66 | 84 |
|
67 | 85 | public Long getSeqNo() { |
68 | 86 | return seqNo; |
69 | 87 | } |
70 | 88 |
|
71 | | - public void setSeqNo(long seqNo) { |
72 | | - this.seqNo = seqNo; |
73 | | - } |
74 | | - |
75 | 89 | public YdbTransaction getTransaction() { |
76 | 90 | return transaction; |
77 | 91 | } |
| 92 | + |
| 93 | + long updateSeqNo(long lastSeqNo) { |
| 94 | + if (seqNo == null) { |
| 95 | + seqNo = lastSeqNo + 1; |
| 96 | + return seqNo; |
| 97 | + } |
| 98 | + return Math.max(lastSeqNo, seqNo); |
| 99 | + } |
| 100 | + |
| 101 | + YdbTopic.StreamWriteMessage.WriteRequest.MessageData toMessageData() { |
| 102 | + return YdbTopic.StreamWriteMessage.WriteRequest.MessageData.newBuilder() |
| 103 | + .setSeqNo(seqNo) |
| 104 | + .setData(UnsafeByteOperations.unsafeWrap(bytes)) |
| 105 | + .setCreatedAt(ProtobufUtils.instantToProto(createdAt)) |
| 106 | + .setUncompressedSize(originLength) |
| 107 | + .addAllMetadataItems(items.stream().map(it -> YdbTopic.MetadataItem.newBuilder() |
| 108 | + .setKey(it.getKey()) |
| 109 | + .setValue(UnsafeByteOperations.unsafeWrap(it.getValue())) |
| 110 | + .build() |
| 111 | + ).collect(Collectors.toList())) |
| 112 | + .build(); |
| 113 | + } |
78 | 114 | } |
0 commit comments