33package com .azure .cosmos .implementation ;
44
55import com .azure .cosmos .ConsistencyLevel ;
6+ import com .azure .cosmos .implementation .directconnectivity .StoreResponse ;
7+ import com .azure .cosmos .implementation .directconnectivity .WFConstants ;
8+ import com .azure .cosmos .implementation .directconnectivity .rntbd .RntbdConstants ;
9+ import com .azure .cosmos .implementation .directconnectivity .rntbd .RntbdFramer ;
610import com .azure .cosmos .implementation .directconnectivity .rntbd .RntbdRequest ;
711import com .azure .cosmos .implementation .directconnectivity .rntbd .RntbdRequestArgs ;
12+ import com .azure .cosmos .implementation .directconnectivity .rntbd .RntbdResponse ;
813import com .azure .cosmos .implementation .http .HttpClient ;
914import com .azure .cosmos .implementation .http .HttpHeaders ;
1015import com .azure .cosmos .implementation .http .HttpRequest ;
11- import com .azure .cosmos .models . FeedRange ;
16+ import com .azure .cosmos .implementation . routing . HexConvert ;
1217import io .netty .buffer .ByteBuf ;
1318import io .netty .buffer .Unpooled ;
1419import io .netty .handler .codec .http .HttpMethod ;
20+ import org .slf4j .Logger ;
21+ import org .slf4j .LoggerFactory ;
1522import reactor .core .publisher .Flux ;
1623import reactor .core .publisher .Mono ;
1724
25+ import java .io .IOException ;
1826import java .net .URI ;
27+ import java .nio .file .Files ;
28+ import java .time .Instant ;
29+ import java .util .Arrays ;
1930import java .util .HashMap ;
31+ import java .util .List ;
2032import java .util .Map ;
33+ import java .util .UUID ;
2134
22- import static com .azure .cosmos .implementation .directconnectivity .WFConstants .BackendHeaders .EFFECTIVE_PARTITION_KEY ;
2335import static com .azure .cosmos .implementation .guava25 .base .Preconditions .checkNotNull ;
2436
2537/**
3042 */
3143public class ThinClientStoreModel extends RxGatewayStoreModel {
3244
45+ private static final Logger logger = LoggerFactory .getLogger (ThinClientStoreModel .class );
46+
47+ private static final List <RntbdConstants .RntbdRequestHeader > thinClientHeadersInOrder = Arrays .asList (
48+ RntbdConstants .RntbdRequestHeader .EffectivePartitionKey ,
49+ RntbdConstants .RntbdRequestHeader .GlobalDatabaseAccountName ,
50+ RntbdConstants .RntbdRequestHeader .DatabaseName ,
51+ RntbdConstants .RntbdRequestHeader .CollectionName ,
52+ RntbdConstants .RntbdRequestHeader .CollectionRid ,
53+ //RntbdConstants.RntbdRequestHeader.ResourceId,
54+ RntbdConstants .RntbdRequestHeader .PayloadPresent ,
55+ RntbdConstants .RntbdRequestHeader .DocumentName ,
56+ RntbdConstants .RntbdRequestHeader .AuthorizationToken ,
57+ RntbdConstants .RntbdRequestHeader .Date );
58+
59+
60+
3361 public ThinClientStoreModel (
3462 DiagnosticsClientContext clientContext ,
3563 ISessionContainer sessionContainer ,
@@ -80,7 +108,37 @@ protected Map<String, String> getDefaultHeaders(
80108 @ Override
81109 public URI getRootUri (RxDocumentServiceRequest request ) {
82110 //var uri = this.globalEndpointManager.resolveServiceEndpoint(request).getThinClientLocationEndpoint();
83- return URI .create ("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/" );
111+ return URI .create ("https://57.155.105.105:10650/" ); // https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/
112+ }
113+
114+ @ Override
115+ public StoreResponse unwrapToStoreResponse (RxDocumentServiceRequest request , int statusCode , HttpHeaders headers , ByteBuf content ) {
116+ if (content == null || content .readableBytes () == 0 ) {
117+ return super .unwrapToStoreResponse (request , statusCode , headers , Unpooled .EMPTY_BUFFER );
118+ }
119+
120+ Instant decodeStartTime = Instant .now ();
121+
122+ if (RntbdFramer .canDecodeHead (content )) {
123+
124+ final RntbdResponse response = RntbdResponse .decode (content );
125+
126+ if (response != null ) {
127+ response .setDecodeEndTime (Instant .now ());
128+ response .setDecodeStartTime (decodeStartTime );
129+
130+ return super .unwrapToStoreResponse (
131+ request ,
132+ response .getStatus ().code (),
133+ new HttpHeaders (response .getHeaders ().asMap (request .getActivityId ())),
134+ response .getContent ()
135+ );
136+ }
137+
138+ return super .unwrapToStoreResponse (request , statusCode , headers , null );
139+ }
140+
141+ throw new IllegalStateException ("Invalid rntbd response" );
84142 }
85143
86144 @ Override
@@ -89,12 +147,14 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque
89147 // todo - neharao1 - validate b/w name() v/s toString()
90148 request .setThinclientHeaders (request .getOperationType ().name (), request .getResourceType ().name ());
91149
92- String epk = request .getPartitionKeyInternal ().getEffectivePartitionKeyString (request .getPartitionKeyInternal (), request .getPartitionKeyDefinition ());
150+ byte [] epk = request .getPartitionKeyInternal ().getEffectivePartitionKeyBytes (request .getPartitionKeyInternal (), request .getPartitionKeyDefinition ());
93151 if (request .properties == null ) {
94152 request .properties = new HashMap <>();
95153 }
96- request .properties .put (EFFECTIVE_PARTITION_KEY , epk );
97- //request.getHeaders().put(EFFECTIVE_PARTITION_KEY, epk);
154+ //request.properties.put(EFFECTIVE_PARTITION_KEY, epk);
155+ //request.properties.put(HttpConstants.HttpHeaders.GLOBAL_DATABASE_ACCOUNT_NAME, "chukangzhongstagesignoff");
156+ request .getHeaders ().put (HttpConstants .HttpHeaders .GLOBAL_DATABASE_ACCOUNT_NAME , "tiagonapoli-cdb-test" ); // "chukangzhongstagesignoff"
157+ request .getHeaders ().put (WFConstants .BackendHeaders .COLLECTION_RID , "cLklAJU8SN0=" );
98158 // todo - neharao1: no concept of a replica / service endpoint that can be passed
99159 RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs (request );
100160
@@ -103,24 +163,46 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque
103163 HttpHeaders headers = this .getHttpHeaders ();
104164
105165 RntbdRequest rntbdRequest = RntbdRequest .from (rntbdRequestArgs );
106-
166+ boolean success = rntbdRequest .setHeaderValue (
167+ RntbdConstants .RntbdRequestHeader .EffectivePartitionKey ,
168+ epk );
169+ if (!success ) {
170+ logger .error ("Failed to update EPK to value {}" , HexConvert .bytesToHex (epk ));
171+ } else {
172+ logger .error ("Updated EPK to value {}" , HexConvert .bytesToHex (epk ));
173+ }
107174 // todo: neharao1 - validate whether Java heap buffer is okay v/s Direct buffer
108175 // todo: eventually need to use pooled buffer
109176 ByteBuf byteBuf = Unpooled .buffer ();
110177
178+ logger .error ("HEADERS: {}" , rntbdRequest .getHeaders ().dumpTokens ());
179+
111180 // todo: lifting the logic from there to encode the RntbdRequest instance into a ByteBuf (ByteBuf is a network compatible format)
112181 // todo: double-check with fabianm to see if RntbdRequest across RNTBD over TCP (Direct connectivity mode) is same as that when using ThinClient proxy
113182 // todo: need to conditionally add some headers (userAgent, replicaId/endpoint, etc)
114- rntbdRequest .encode (byteBuf );
183+ rntbdRequest .encode (byteBuf , true );
184+
185+ byte [] contentAsByteArray = new byte [byteBuf .writerIndex ()];
186+ byteBuf .getBytes (0 , contentAsByteArray , 0 , byteBuf .writerIndex ());
187+
188+ try {
189+ Files .write (java .nio .file .Paths .get ("E:\\ Temp\\ java" + UUID .randomUUID () + ".bin" ), contentAsByteArray );
190+ } catch (IOException e ) {
191+ e .printStackTrace ();
192+ }
115193
116194 return new HttpRequest (
117195 HttpMethod .POST ,
118196 //requestUri,
119- URI .create ("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/" ),
197+ //https://thinclient-performancetests-eastus2.documents-staging.windows-ppe.net:10650
198+ //https://cdb-ms-stage-eastus2-fe2-sql.eastus2.cloudapp.azure.com:10650
199+ //https://57.155.105.105:10650/
200+ // https://tiagonapoli-cdb-test-westus3.documents.azure.com:10650
201+ URI .create ("https://57.155.105.105:10650/" ), // https://127.0.0.1:10650/ //https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/ // thinclient-performancetests-eastus2.documents-staging.windows-ppe.net cdb-ms-stage-eastus2-fe2-sql.eastus2.cloudapp.azure.com
120202 //requestUri.getPort(),
121203 10650 ,
122204 headers ,
123- Flux .just (byteBuf . array () ));
205+ Flux .just (contentAsByteArray ));
124206 }
125207
126208 private HttpHeaders getHttpHeaders () {
0 commit comments