diff --git a/client/rest-http-client/build.gradle b/client/rest-http-client/build.gradle new file mode 100644 index 0000000000000..69a6813f0329e --- /dev/null +++ b/client/rest-http-client/build.gradle @@ -0,0 +1,134 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +import de.thetaphi.forbiddenapis.gradle.CheckForbiddenApis + +apply plugin: 'opensearch.build' +apply plugin: 'opensearch.publish' +apply from: "$rootDir/gradle/fips.gradle" + +java { + targetCompatibility = JavaVersion.VERSION_21 + sourceCompatibility = JavaVersion.VERSION_21 +} + +base { + group = 'org.opensearch.client' + archivesName = 'opensearch-rest-http-client' +} + +dependencies { + api "commons-codec:commons-codec:${versions.commonscodec}" + api "commons-logging:commons-logging:${versions.commonslogging}" + api "org.slf4j:slf4j-api:${versions.slf4j}" + fipsRuntimeOnly "org.bouncycastle:bc-fips:${versions.bouncycastle_jce}" + fipsRuntimeOnly "org.bouncycastle:bctls-fips:${versions.bouncycastle_tls}" + fipsRuntimeOnly "org.bouncycastle:bcutil-fips:${versions.bouncycastle_util}" + + // reactor + api "io.projectreactor:reactor-core:${versions.reactor}" + api "org.reactivestreams:reactive-streams:${versions.reactivestreams}" + + testImplementation "com.carrotsearch.randomizedtesting:randomizedtesting-runner:${versions.randomizedrunner}" + testImplementation "junit:junit:${versions.junit}" + testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}" + testImplementation "org.mockito:mockito-core:${versions.mockito}" + testImplementation "org.objenesis:objenesis:${versions.objenesis}" + testImplementation "net.bytebuddy:byte-buddy:${versions.bytebuddy}" + testImplementation "net.bytebuddy:byte-buddy-agent:${versions.bytebuddy}" + testImplementation "org.apache.logging.log4j:log4j-api:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-core:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-jul:${versions.log4j}" + testImplementation "org.apache.logging.log4j:log4j-slf4j2-impl:${versions.log4j}" + testImplementation "io.projectreactor:reactor-test:${versions.reactor}" +} + +tasks.named("dependencyLicenses").configure { + mapping from: /bc.*/, to: 'bouncycastle' +} + +tasks.withType(CheckForbiddenApis).configureEach { + //client does not depend on server, so only jdk and http signatures should be checked + replaceSignatureFiles('jdk-signatures') +} + +tasks.named('forbiddenApisTest').configure { + //we are using jdk-internal instead of jdk-non-portable to allow for com.sun.net.httpserver.* usage + bundledSignatures -= 'jdk-non-portable' + bundledSignatures += 'jdk-internal' +} + +testingConventions { + naming.clear() + naming { + Tests { + baseClass 'org.opensearch.internal.httpclient.RestHttpClientTestCase' + } + } +} + +thirdPartyAudit { + ignoreMissingClasses( + //commons-logging optional dependencies + 'org.apache.avalon.framework.logger.Logger', + 'org.apache.log.Hierarchy', + 'org.apache.log.Logger', + 'org.apache.log4j.Level', + 'org.apache.log4j.Logger', + 'org.apache.log4j.Priority', + 'org.apache.logging.log4j.Level', + 'org.apache.logging.log4j.LogManager', + 'org.apache.logging.log4j.Marker', + 'org.apache.logging.log4j.MarkerManager', + 'org.apache.logging.log4j.spi.AbstractLoggerAdapter', + 'org.apache.logging.log4j.spi.ExtendedLogger', + 'org.apache.logging.log4j.spi.LoggerAdapter', + 'org.apache.logging.log4j.spi.LoggerContext', + 'org.apache.logging.log4j.spi.LoggerContextFactory', + 'org.apache.logging.log4j.util.StackLocatorUtil', + //commons-logging provided dependencies + 'javax.servlet.ServletContextEvent', + 'javax.servlet.ServletContextListener', + 'io.micrometer.context.ContextAccessor', + 'io.micrometer.context.ContextRegistry', + 'io.micrometer.context.ContextSnapshot', + 'io.micrometer.context.ContextSnapshot$Scope', + 'io.micrometer.context.ContextSnapshotFactory', + 'io.micrometer.context.ContextSnapshotFactory$Builder', + 'io.micrometer.context.ThreadLocalAccessor', + 'io.micrometer.core.instrument.Clock', + 'io.micrometer.core.instrument.Counter', + 'io.micrometer.core.instrument.Counter$Builder', + 'io.micrometer.core.instrument.DistributionSummary', + 'io.micrometer.core.instrument.DistributionSummary$Builder', + 'io.micrometer.core.instrument.Meter', + 'io.micrometer.core.instrument.MeterRegistry', + 'io.micrometer.core.instrument.Metrics', + 'io.micrometer.core.instrument.Tag', + 'io.micrometer.core.instrument.Tags', + 'io.micrometer.core.instrument.Timer', + 'io.micrometer.core.instrument.Timer$Builder', + 'io.micrometer.core.instrument.Timer$Sample', + 'io.micrometer.core.instrument.binder.jvm.ExecutorServiceMetrics', + 'io.micrometer.core.instrument.composite.CompositeMeterRegistry', + 'io.micrometer.core.instrument.search.Search', + 'reactor.blockhound.BlockHound$Builder', + 'reactor.blockhound.integration.BlockHoundIntegration' + ) + ignoreViolations( + 'reactor.core.publisher.CallSiteSupplierFactory$SharedSecretsCallSiteSupplierFactory', + 'reactor.core.publisher.CallSiteSupplierFactory$SharedSecretsCallSiteSupplierFactory$TracingException' + ) +} + +tasks.named("missingJavadoc").configure { + it.enabled = false +} diff --git a/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 b/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 new file mode 100644 index 0000000000000..bd2f333cb12d0 --- /dev/null +++ b/client/rest-http-client/licenses/bc-fips-2.1.2.jar.sha1 @@ -0,0 +1 @@ +061fbe8383f70489dda95a11a2a4739eb818ff2c \ No newline at end of file diff --git a/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 b/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 new file mode 100644 index 0000000000000..7266ec5abf10a --- /dev/null +++ b/client/rest-http-client/licenses/bctls-fips-2.1.20.jar.sha1 @@ -0,0 +1 @@ +9c0632a6c5ca09a86434cf5e02e72c221e1c930f \ No newline at end of file diff --git a/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 b/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 new file mode 100644 index 0000000000000..73b19722430fb --- /dev/null +++ b/client/rest-http-client/licenses/bcutil-fips-2.1.4.jar.sha1 @@ -0,0 +1 @@ +1d37b7a28560684f5b8e4fd65478c9130d4015d0 \ No newline at end of file diff --git a/client/rest-http-client/licenses/bouncycastle-LICENSE.txt b/client/rest-http-client/licenses/bouncycastle-LICENSE.txt new file mode 100644 index 0000000000000..5c7c14696849d --- /dev/null +++ b/client/rest-http-client/licenses/bouncycastle-LICENSE.txt @@ -0,0 +1,14 @@ +Copyright (c) 2000 - 2023 The Legion of the Bouncy Castle Inc. (https://www.bouncycastle.org) + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, +and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the +Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE +WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/client/rest-http-client/licenses/bouncycastle-NOTICE.txt b/client/rest-http-client/licenses/bouncycastle-NOTICE.txt new file mode 100644 index 0000000000000..8b137891791fe --- /dev/null +++ b/client/rest-http-client/licenses/bouncycastle-NOTICE.txt @@ -0,0 +1 @@ + diff --git a/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 b/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 new file mode 100644 index 0000000000000..01a6a8f446302 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-1.18.0.jar.sha1 @@ -0,0 +1 @@ +ee45d1cf6ec2cc2b809ff04b4dc7aec858e0df8f \ No newline at end of file diff --git a/client/rest-http-client/licenses/commons-codec-LICENSE.txt b/client/rest-http-client/licenses/commons-codec-LICENSE.txt new file mode 100644 index 0000000000000..d645695673349 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/client/rest-http-client/licenses/commons-codec-NOTICE.txt b/client/rest-http-client/licenses/commons-codec-NOTICE.txt new file mode 100644 index 0000000000000..1da9af50f6008 --- /dev/null +++ b/client/rest-http-client/licenses/commons-codec-NOTICE.txt @@ -0,0 +1,17 @@ +Apache Commons Codec +Copyright 2002-2014 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + +src/test/org/apache/commons/codec/language/DoubleMetaphoneTest.java +contains test data from http://aspell.net/test/orig/batch0.tab. +Copyright (C) 2002 Kevin Atkinson (kevina@gnu.org) + +=============================================================================== + +The content of package org.apache.commons.codec.language.bm has been translated +from the original php source code available at http://stevemorse.org/phoneticinfo.htm +with permission from the original authors. +Original source copyright: +Copyright (c) 2008 Alexander Beider & Stephen P. Morse. diff --git a/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 b/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 new file mode 100644 index 0000000000000..f7ddad61aaeaa --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-1.3.5.jar.sha1 @@ -0,0 +1 @@ +a3fcc5d3c29b2b03433aa2d2f2d2c1b1638924a1 \ No newline at end of file diff --git a/client/rest-http-client/licenses/commons-logging-LICENSE.txt b/client/rest-http-client/licenses/commons-logging-LICENSE.txt new file mode 100644 index 0000000000000..d645695673349 --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/client/rest-http-client/licenses/commons-logging-NOTICE.txt b/client/rest-http-client/licenses/commons-logging-NOTICE.txt new file mode 100644 index 0000000000000..a37977d45a168 --- /dev/null +++ b/client/rest-http-client/licenses/commons-logging-NOTICE.txt @@ -0,0 +1,6 @@ +Apache Commons Logging +Copyright 2003-2013 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). + diff --git a/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 b/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 new file mode 100644 index 0000000000000..45a80e3f7e361 --- /dev/null +++ b/client/rest-http-client/licenses/reactive-streams-1.0.4.jar.sha1 @@ -0,0 +1 @@ +3864a1320d97d7b045f729a326e1e077661f31b7 \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactive-streams-LICENSE.txt b/client/rest-http-client/licenses/reactive-streams-LICENSE.txt new file mode 100644 index 0000000000000..1e3c7e7c77495 --- /dev/null +++ b/client/rest-http-client/licenses/reactive-streams-LICENSE.txt @@ -0,0 +1,21 @@ +MIT No Attribution + +Copyright 2014 Reactive Streams + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactive-streams-NOTICE.txt b/client/rest-http-client/licenses/reactive-streams-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 b/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 new file mode 100644 index 0000000000000..6cdefa9580bdb --- /dev/null +++ b/client/rest-http-client/licenses/reactor-core-3.8.6.jar.sha1 @@ -0,0 +1 @@ +76285d63d5da4ed8679357628dd5309b0feb77ff \ No newline at end of file diff --git a/client/rest-http-client/licenses/reactor-core-LICENSE.txt b/client/rest-http-client/licenses/reactor-core-LICENSE.txt new file mode 100644 index 0000000000000..e5583c184e67a --- /dev/null +++ b/client/rest-http-client/licenses/reactor-core-LICENSE.txt @@ -0,0 +1,201 @@ +Apache License + Version 2.0, January 2004 + https://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright {yyyy} {name of copyright owner} + + 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 + + https://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. diff --git a/client/rest-http-client/licenses/reactor-core-NOTICE.txt b/client/rest-http-client/licenses/reactor-core-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 b/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 new file mode 100644 index 0000000000000..435f6c13a28b6 --- /dev/null +++ b/client/rest-http-client/licenses/slf4j-api-2.0.17.jar.sha1 @@ -0,0 +1 @@ +d9e58ac9c7779ba3bf8142aff6c830617a7fe60f \ No newline at end of file diff --git a/client/rest-http-client/licenses/slf4j-api-LICENSE.txt b/client/rest-http-client/licenses/slf4j-api-LICENSE.txt new file mode 100644 index 0000000000000..8fda22f4d72f6 --- /dev/null +++ b/client/rest-http-client/licenses/slf4j-api-LICENSE.txt @@ -0,0 +1,21 @@ +Copyright (c) 2004-2014 QOS.ch +All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/client/rest-http-client/licenses/slf4j-api-NOTICE.txt b/client/rest-http-client/licenses/slf4j-api-NOTICE.txt new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java new file mode 100644 index 0000000000000..aecc1345100c3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/AsyncResponseProducer.java @@ -0,0 +1,42 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Vector; +import java.util.concurrent.Flow.Subscriber; +import java.util.concurrent.Flow.Subscription; + +class AsyncResponseProducer implements Subscriber> { + private Subscription subscription; + private final List buffers = new Vector<>(); + + @Override + public void onSubscribe(Subscription subscription) { + this.subscription = subscription; + subscription.request(1); + } + + @Override + public void onNext(List item) { + buffers.addAll(item); + subscription.request(1); + } + + @Override + public void onError(Throwable throwable) {} + + @Override + public void onComplete() {} + + public List getResult() { + return buffers; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java new file mode 100644 index 0000000000000..4b1ef84f8a702 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/BodyUtils.java @@ -0,0 +1,268 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.io.UncheckedIOException; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.nio.channels.Channels; +import java.nio.channels.WritableByteChannel; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Flow; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Mono; + +/** + * Helper class to deal with request / response bodies. + */ +final class BodyUtils { + static Mono getBody(HttpRequest httpRequest) { + return httpRequest.bodyPublisher().map(JdkFlowAdapter::flowPublisherToFlux).map(Mono::from).orElseGet(Mono::empty); + } + + static String getBodyAsString(Response response) { + return getBodyAsString(response.entity()); + } + + static Mono getBodyAsString(HttpRequest httpRequest) { + return httpRequest.bodyPublisher().map(p -> { + var bodySubscriber = HttpResponse.BodySubscribers.ofString(StandardCharsets.UTF_8); + var flowSubscriber = new StringSubscriber(bodySubscriber); + p.subscribe(flowSubscriber); + return Mono.fromCompletionStage(bodySubscriber.getBody()); + }).orElseGet(Mono::empty); + } + + static String getBodyAsString(List body) { + final StringBuilder builder = new StringBuilder(); + if (body != null && body.isEmpty() == false) { + for (ByteBuffer chunk : body) { + chunk.mark(); + builder.append(StandardCharsets.UTF_8.decode(chunk).toString()); + chunk.reset(); + } + } + return builder.toString(); + } + + static List compress(List body) { + if (body == null || body.isEmpty()) { + return body; + } else { + body.stream().forEach(ByteBuffer::mark); + try (ByteBufferOutputStream bbout = new ByteBufferOutputStream(); OutputStream out = new GZIPOutputStream(bbout)) { + try (WritableByteChannel channel = Channels.newChannel(out)) { + for (ByteBuffer buffer : body) { + channel.write(buffer); + } + out.flush(); + } + return bbout.getBufferList(); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.stream().forEach(ByteBuffer::reset); + } + } + } + + static ByteBuffer compress(ByteBuffer body) { + if (body == null) { + return body; + } else { + body.mark(); + try (ByteBufferOutputStream bbout = new ByteBufferOutputStream(); OutputStream out = new GZIPOutputStream(bbout)) { + try (WritableByteChannel channel = Channels.newChannel(out)) { + channel.write(body); + out.flush(); + } + + final List bufferList = bbout.getBufferList(); + if (bufferList.isEmpty() == false) { + return bufferList.get(0); + } else { + // We should never end up here + return ByteBuffer.allocate(0); /* empty byte buffer */ + } + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.reset(); + } + } + } + + static ByteBuffer decompress(ByteBuffer body) { + if (body == null) { + return body; + } else { + body.mark(); + + try (ByteBufferInputStream bbin = new ByteBufferInputStream(List.of(body)); InputStream in = new GZIPInputStream(bbin)) { + return ByteBuffer.wrap(in.readAllBytes()); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.reset(); + } + } + } + + static List decompress(List body) { + if (body == null || body.isEmpty()) { + return body; + } else { + body.stream().forEach(ByteBuffer::mark); + try (ByteBufferInputStream bbin = new ByteBufferInputStream(body); InputStream in = new GZIPInputStream(bbin)) { + return List.of(ByteBuffer.wrap(in.readAllBytes())); + } catch (final IOException ex) { + throw new UncheckedIOException(ex); + } finally { + body.stream().forEach(ByteBuffer::reset); + } + } + } + + /** + * See please https://github.com/justinsb/avro/blob/master/src/java/org/apache/avro/ipc/ByteBufferOutputStream.java + */ + private final static class ByteBufferOutputStream extends OutputStream { + private static final int BUFFER_SIZE = 8192; + private List buffers; + + ByteBufferOutputStream() { + reset(); + } + + /** Returns all data written and resets the stream to be empty. */ + List getBufferList() { + List result = buffers; + reset(); + for (ByteBuffer buffer : result) { + buffer.flip(); + } + return result; + } + + private void reset() { + buffers = new ArrayList(1); + buffers.add(ByteBuffer.allocate(BUFFER_SIZE)); + } + + @Override + public void write(int b) { + ByteBuffer buffer = buffers.get(buffers.size() - 1); + if (buffer.remaining() < 1) { + buffer = ByteBuffer.allocate(BUFFER_SIZE); + buffers.add(buffer); + } + buffer.put((byte) b); + } + + @Override + public void write(byte[] b, int off, int len) { + ByteBuffer buffer = buffers.get(buffers.size() - 1); + int remaining = buffer.remaining(); + while (len > remaining) { + buffer.put(b, off, remaining); + len -= remaining; + off += remaining; + buffer = ByteBuffer.allocate(BUFFER_SIZE); + buffers.add(buffer); + remaining = buffer.remaining(); + } + buffer.put(b, off, len); + } + } + + /** + * See please https://github.com/justinsb/avro/blob/master/src/java/org/apache/avro/ipc/ByteBufferInputStream.java + */ + private static final class ByteBufferInputStream extends InputStream { + private List buffers; + private int current; + + ByteBufferInputStream(List buffers) { + this.buffers = buffers; + } + + /** @see InputStream#read() + * @throws EOFException if EOF is reached. */ + @Override + public int read() throws IOException { + return getBuffer().get() & 0xff; + } + + /** @see InputStream#read(byte[], int, int) + * @throws EOFException if EOF is reached before reading all the bytes. */ + @Override + public int read(byte[] b, int off, int len) throws IOException { + if (len == 0) return 0; + ByteBuffer buffer = getBuffer(); + int remaining = buffer.remaining(); + if (len > remaining) { + buffer.get(b, off, remaining); + return remaining; + } else { + buffer.get(b, off, len); + return len; + } + } + + /** Returns the next non-empty buffer. + * @throws EOFException if EOF is reached before reading all the bytes. + */ + private ByteBuffer getBuffer() throws IOException { + while (current < buffers.size()) { + ByteBuffer buffer = buffers.get(current); + if (buffer.hasRemaining()) return buffer; + current++; + } + throw new EOFException(); + } + } + + private static final class StringSubscriber implements Flow.Subscriber { + final HttpResponse.BodySubscriber delegate; + + StringSubscriber(HttpResponse.BodySubscriber delegate) { + this.delegate = delegate; + } + + @Override + public void onSubscribe(Flow.Subscription subscription) { + delegate.onSubscribe(subscription); + } + + @Override + public void onNext(ByteBuffer item) { + delegate.onNext(List.of(item)); + } + + @Override + public void onError(Throwable throwable) { + delegate.onError(throwable); + } + + @Override + public void onComplete() { + delegate.onComplete(); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java new file mode 100644 index 0000000000000..6357ac163e410 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Cancellable.java @@ -0,0 +1,120 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.Future; + +/** + * Represents an operation that can be cancelled. + * Returned when executing async requests through {@link RestHttpClient#performRequestAsync(Request, ResponseListener)}, so that the request + * can be cancelled if needed. Cancelling a request will result in calling {@link Future#cancel(boolean mayInterruptIfRunning)} on the + * underlying request future object. Note that cancelling a request does not automatically translate to aborting its execution on the + * server side, which needs to be specifically implemented in each API. + */ +public class Cancellable { + static final Cancellable NO_OP = new Cancellable(null) { + @Override + public boolean cancel() { + throw new UnsupportedOperationException(); + } + + @Override + void runIfNotCancelled(Runnable runnable) { + throw new UnsupportedOperationException(); + } + }; + + static Cancellable fromFuture(Future f) { + return new Cancellable(f); + } + + private final Future future; + + private Cancellable(Future f) { + this.future = f; + } + + /** + * Cancels the on-going request that is associated with the current instance of {@link Cancellable}. + * + */ + public synchronized boolean cancel() { + return this.future.cancel(true); + } + + /** + * Executes some arbitrary code if the on-going request has not been cancelled, otherwise throws {@link CancellationException}. + * This is needed to guarantee that cancelling a request works correctly even in case {@link #cancel()} is called between different + * attempts of the same request. + * If the request has already been cancelled we don't go ahead with the next attempt, and artificially raise the + * {@link CancellationException}, otherwise we run the provided {@link Runnable} which will reset the request and send the next attempt. + * Note that this method must be synchronized as well as the {@link #cancel()} method, to prevent a request from being cancelled + * when there is no future to cancel, which would make cancelling the request a no-op. + */ + synchronized void runIfNotCancelled(Runnable runnable) { + if (this.future.isCancelled()) { + throw newCancellationException(); + } + runnable.run(); + } + + /** + * Executes some arbitrary code if the on-going request has not been cancelled, otherwise throws {@link CancellationException}. + * This is needed to guarantee that cancelling a request works correctly even in case {@link #cancel()} is called between different + * attempts of the same request. The {@link #cancel()} method can be called at anytime, + * and we need to handle the case where it gets called while there is no request being executed as one attempt may have failed and + * the subsequent attempt has not been started yet. + * If the request has already been cancelled we don't go ahead with the next attempt, and artificially raise the + * {@link CancellationException}, otherwise we run the provided {@link Runnable} which will reset the request and send the next attempt. + * Note that this method must be synchronized as well as the {@link #cancel()} method, to prevent a request from being cancelled + * when there is no future to cancel, which would make cancelling the request a no-op. + */ + synchronized T callIfNotCancelled(Callable callable) throws IOException { + if (this.future.isCancelled()) { + throw newCancellationException(); + } + try { + return callable.call(); + } catch (final IOException ex) { + throw ex; + } catch (final Exception ex) { + throw new IOException(ex); + } + } + + static CancellationException newCancellationException() { + return new CancellationException("request was cancelled"); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java new file mode 100644 index 0000000000000..4060891dd48dd --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/CompressedResponse.java @@ -0,0 +1,41 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.InputStream; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Objects; + +final record CompressedResponse(RequestLine requestLine, HttpHost host, HttpResponse httpResponse, List entity) + implements + Response { + CompressedResponse { + requestLine = Objects.requireNonNull(requestLine, "requestLine cannot be null"); + host = Objects.requireNonNull(host, "host cannot be null"); + httpResponse = Objects.requireNonNull(httpResponse, "response cannot be null"); + } + + /** + * Returns the response body available, null otherwise + * @see InputStream + */ + public List entity() { + return BodyUtils.decompress(entity); + } + + /** + * Convert response to string representation + */ + @Override + public String toString() { + return "Response{" + "requestLine=" + requestLine() + ", host=" + host() + ", response=" + statusLine() + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java new file mode 100644 index 0000000000000..de2e5161ea9c9 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/DeadHostState.java @@ -0,0 +1,125 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +/** + * Holds the state of a dead connection to a host. Keeps track of how many failed attempts were performed and + * when the host should be retried (based on number of previous failed attempts). + * Class is immutable, a new copy of it should be created each time the state has to be changed. + */ +final class DeadHostState implements Comparable { + + private static final long MIN_CONNECTION_TIMEOUT_NANOS = TimeUnit.MINUTES.toNanos(1); + static final long MAX_CONNECTION_TIMEOUT_NANOS = TimeUnit.MINUTES.toNanos(30); + static final Supplier DEFAULT_TIME_SUPPLIER = System::nanoTime; + + private final int failedAttempts; + private final long deadUntilNanos; + private final Supplier timeSupplier; + + /** + * Build the initial dead state of a host. Useful when a working host stops functioning + * and needs to be marked dead after its first failure. In such case the host will be retried after a minute or so. + * + * @param timeSupplier a way to supply the current time and allow for unit testing + */ + DeadHostState(Supplier timeSupplier) { + this.failedAttempts = 1; + this.deadUntilNanos = timeSupplier.get() + MIN_CONNECTION_TIMEOUT_NANOS; + this.timeSupplier = timeSupplier; + } + + /** + * Build the dead state of a host given its previous dead state. Useful when a host has been failing before, hence + * it already failed for one or more consecutive times. The more failed attempts we register the longer we wait + * to retry that same host again. Minimum is 1 minute (for a node the only failed once created + * through {@link #DeadHostState(Supplier)}), maximum is 30 minutes (for a node that failed more than 10 consecutive times) + * + * @param previousDeadHostState the previous state of the host which allows us to increase the wait till the next retry attempt + */ + DeadHostState(DeadHostState previousDeadHostState) { + long timeoutNanos = (long) Math.min( + MIN_CONNECTION_TIMEOUT_NANOS * 2 * Math.pow(2, previousDeadHostState.failedAttempts * 0.5 - 1), + MAX_CONNECTION_TIMEOUT_NANOS + ); + this.deadUntilNanos = previousDeadHostState.timeSupplier.get() + timeoutNanos; + this.failedAttempts = previousDeadHostState.failedAttempts + 1; + this.timeSupplier = previousDeadHostState.timeSupplier; + } + + /** + * Indicates whether it's time to retry to failed host or not. + * + * @return true if the host should be retried, false otherwise + */ + boolean shallBeRetried() { + return timeSupplier.get() - deadUntilNanos > 0; + } + + /** + * Returns the timestamp (nanos) till the host is supposed to stay dead without being retried. + * After that the host should be retried. + */ + long getDeadUntilNanos() { + return deadUntilNanos; + } + + int getFailedAttempts() { + return failedAttempts; + } + + @Override + public int compareTo(DeadHostState other) { + if (timeSupplier != other.timeSupplier) { + throw new IllegalArgumentException( + "can't compare DeadHostStates holding different time suppliers as they may " + "be based on different clocks" + ); + } + return Long.compare(deadUntilNanos, other.deadUntilNanos); + } + + @Override + public String toString() { + return "DeadHostState{" + + "failedAttempts=" + + failedAttempts + + ", deadUntilNanos=" + + deadUntilNanos + + ", timeSupplier=" + + timeSupplier + + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java new file mode 100644 index 0000000000000..cb7d7d6106a20 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/HttpHost.java @@ -0,0 +1,52 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Objects; + +/** + * Represents HTTP host (scheme, hostname and port) + * Note: This is an experimental API. + */ +public record HttpHost(String scheme, String hostname, int port) { + /** + * Converts the host to string + */ + @Override + public String toString() { + final StringBuilder buffer = new StringBuilder(); + buffer.append(this.scheme); + buffer.append("://"); + buffer.append(this.hostname); + if (this.port != -1) { + buffer.append(':'); + buffer.append(Integer.toString(this.port)); + } + return buffer.toString(); + } + + public static HttpHost create(String uriStr) throws URISyntaxException { + Objects.requireNonNull(uriStr); + + String text = uriStr; + String scheme = null; + final int schemeIdx = text.indexOf("://"); + if (schemeIdx > 0) { + scheme = text.substring(0, schemeIdx); + if (scheme.isBlank()) { + throw new URISyntaxException(uriStr, "scheme contains blanks"); + } + } + + final URI uri = new URI(uriStr); + return new HttpHost(scheme, uri.getHost(), uri.getPort()); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java new file mode 100644 index 0000000000000..c892b6fae2b64 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Node.java @@ -0,0 +1,278 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.TreeSet; + +/** + * Metadata about an {@link HttpHost} running OpenSearch. + */ +public class Node { + /** + * Address that this host claims is its primary contact point. + */ + private final HttpHost host; + /** + * Addresses on which the host is listening. These are useful to have + * around because they allow you to find a host based on any address it + * is listening on. + */ + private final Set boundHosts; + /** + * Name of the node as configured by the {@code node.name} attribute. + */ + private final String name; + /** + * Version of OpenSearch that the node is running or {@code null} + * if we don't know the version. + */ + private final String version; + /** + * Roles that the OpenSearch process on the host has or {@code null} + * if we don't know what roles the node has. + */ + private final Roles roles; + /** + * Attributes declared on the node. + */ + private final Map> attributes; + + /** + * Create a {@linkplain Node} with metadata. All parameters except + * {@code host} are nullable and implementations of {@link NodeSelector} + * need to decide what to do in their absence. + * + * @param host primary host address + * @param boundHosts addresses on which the host is listening + * @param name name of the node + * @param version version of OpenSearch + * @param roles roles that the OpenSearch process has on the host + * @param attributes attributes declared on the node + */ + public Node(HttpHost host, Set boundHosts, String name, String version, Roles roles, Map> attributes) { + if (host == null) { + throw new IllegalArgumentException("host cannot be null"); + } + this.host = host; + this.boundHosts = boundHosts; + this.name = name; + this.version = version; + this.roles = roles; + this.attributes = attributes; + } + + /** + * Create a {@linkplain Node} without any metadata. + * + * @param host primary host address + */ + public Node(HttpHost host) { + this(host, null, null, null, null, null); + } + + /** + * Contact information for the host. + */ + public HttpHost getHost() { + return host; + } + + /** + * Addresses on which the host is listening. These are useful to have + * around because they allow you to find a host based on any address it + * is listening on. + */ + public Set getBoundHosts() { + return boundHosts; + } + + /** + * The {@code node.name} of the node. + */ + public String getName() { + return name; + } + + /** + * Version of OpenSearch that the node is running or {@code null} + * if we don't know the version. + */ + public String getVersion() { + return version; + } + + /** + * Roles that the OpenSearch process on the host has or {@code null} + * if we don't know what roles the node has. + */ + public Roles getRoles() { + return roles; + } + + /** + * Attributes declared on the node. + */ + public Map> getAttributes() { + return attributes; + } + + /** + * Convert node to string representation + */ + @Override + public String toString() { + StringBuilder b = new StringBuilder(); + b.append("[host=").append(host); + if (boundHosts != null) { + b.append(", bound=").append(boundHosts); + } + if (name != null) { + b.append(", name=").append(name); + } + if (version != null) { + b.append(", version=").append(version); + } + if (roles != null) { + b.append(", roles=").append(roles); + } + if (attributes != null) { + b.append(", attributes=").append(attributes); + } + return b.append(']').toString(); + } + + /** + * Compare two nodes for equality + * @param obj node instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || obj.getClass() != getClass()) { + return false; + } + Node other = (Node) obj; + return host.equals(other.host) + && Objects.equals(boundHosts, other.boundHosts) + && Objects.equals(name, other.name) + && Objects.equals(version, other.version) + && Objects.equals(roles, other.roles) + && Objects.equals(attributes, other.attributes); + } + + /** + * Calculate the hash code of the node + */ + @Override + public int hashCode() { + return Objects.hash(host, boundHosts, name, version, roles, attributes); + } + + /** + * Role information about an OpenSearch process. + */ + public static final class Roles { + + private final Set roles; + + /** + * Create a {@link Roles} instance of the given string set. + * + * @param roles set of role names. + */ + public Roles(final Set roles) { + this.roles = new TreeSet<>(roles); + } + + /** + * Returns whether or not the node could be elected cluster-manager. + */ + public boolean isClusterManagerEligible() { + return roles.contains("master") || roles.contains("cluster_manager"); + } + + /** + * Returns whether or not the node stores data. + */ + public boolean isData() { + return roles.contains("data"); + } + + /** + * Returns whether or not the node runs ingest pipelines. + */ + public boolean isIngest() { + return roles.contains("ingest"); + } + + /** + * Returns whether the node is dedicated to provide search capability. + */ + public boolean isSearch() { + return roles.contains("search"); + } + + /** + * Convert roles to string representation + */ + @Override + public String toString() { + return String.join(",", roles); + } + + /** + * Compare two roles for equality + * @param obj roles instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || obj.getClass() != getClass()) { + return false; + } + Roles other = (Roles) obj; + return roles.equals(other.roles); + } + + /** + * Calculate the hash code of the roles + */ + @Override + public int hashCode() { + return roles.hashCode(); + } + + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java new file mode 100644 index 0000000000000..4a4856a920b48 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NodeSelector.java @@ -0,0 +1,105 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.Iterator; + +/** + * Selects nodes that can receive requests. Used to keep requests away + * from cluster-manager nodes or to send them to nodes with a particular attribute. + * Use with {@link RestHttpClientBuilder#setNodeSelector(NodeSelector)}. + */ +public interface NodeSelector { + /** + * Select the {@link Node}s to which to send requests. This is called with + * a mutable {@link Iterable} of {@linkplain Node}s in the order that the + * rest client would prefer to use them and implementers should remove + * nodes from the that should not receive the request. Implementers may + * iterate the nodes as many times as they need. + *

+ * This may be called twice per request: first for "living" nodes that + * have not been denylisted by previous errors. If the selector removes + * all nodes from the list or if there aren't any living nodes then the + * {@link RestHttpClient} will call this method with a list of "dead" nodes. + *

+ * Implementers should not rely on the ordering of the nodes. + * + * @param nodes the {@link Node}s targeted for the sending requests + */ + void select(Iterable nodes); + /* + * We were fairly careful with our choice of Iterable here. The caller has + * a List but reordering the list is likely to break round robin. Luckily + * Iterable doesn't allow any reordering. + */ + + /** + * Selector that matches any node. + */ + NodeSelector ANY = new NodeSelector() { + @Override + public void select(Iterable nodes) { + // Intentionally does nothing + } + + @Override + public String toString() { + return "ANY"; + } + }; + + /** + * Selector that matches any node that has metadata and doesn't + * have the {@code cluster_manager} role OR it has the data {@code data} + * role. + */ + NodeSelector SKIP_DEDICATED_CLUSTER_MANAGERS = new NodeSelector() { + @Override + public void select(Iterable nodes) { + for (Iterator itr = nodes.iterator(); itr.hasNext();) { + Node node = itr.next(); + if (node.getRoles() == null) continue; + if (node.getRoles().isClusterManagerEligible() + && false == node.getRoles().isData() + && false == node.getRoles().isIngest()) { + itr.remove(); + } + } + } + + @Override + public String toString() { + return "SKIP_DEDICATED_CLUSTER_MANAGERS"; + } + }; +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java new file mode 100644 index 0000000000000..67137c468d2d6 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/NonCompressedResponse.java @@ -0,0 +1,32 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Objects; + +final record NonCompressedResponse(RequestLine requestLine, HttpHost host, HttpResponse httpResponse, List entity) + implements + Response { + NonCompressedResponse { + requestLine = Objects.requireNonNull(requestLine, "requestLine cannot be null"); + host = Objects.requireNonNull(host, "host cannot be null"); + httpResponse = Objects.requireNonNull(httpResponse, "response cannot be null"); + } + + /** + * Convert response to string representation + */ + @Override + public String toString() { + return "Response{" + "requestLine=" + requestLine() + ", host=" + host() + ", response=" + statusLine() + '}'; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java new file mode 100644 index 0000000000000..ed4cc1e8f19d4 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Request.java @@ -0,0 +1,160 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import static java.util.Collections.unmodifiableMap; + +/** + * HTTP Request to OpenSearch. + * Note: This is an experimental API. + */ +public record Request(String method, String endpoint, Map parameters, BodyPublisher entity, RequestOptions options) { + + public Request { + method = Objects.requireNonNull(method, "method cannot be null"); + endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + parameters = parameters == null ? new HashMap<>() : parameters; + options = options == null ? RequestOptions.DEFAULT : options; + } + + /** + * Query string parameters. The returned map is an unmodifiable view of the + * map in the request so calls to {@link Map#put(Object, Object)} + * will change it. + */ + @Override + public Map parameters() { + if (options.getParameters().isEmpty()) { + return unmodifiableMap(parameters); + } else { + Map combinedParameters = new HashMap<>(parameters); + combinedParameters.putAll(options.getParameters()); + return unmodifiableMap(combinedParameters); + } + } + + public static Request.Builder newRequest(String method, String endpoint) { + return new Request.Builder(method, endpoint); + } + + public static final class Builder { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private BodyPublisher entity; + private RequestOptions options = RequestOptions.DEFAULT; + + private Builder(String method, String endpoint) { + this.method = Objects.requireNonNull(method, "method cannot be null"); + this.endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public void setOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + } + + /** + * Add a query string parameter. + * @param name the name of the url parameter. Must not be null. + * @param value the value of the url url parameter. If {@code null} then + * the parameter is sent as {@code name} rather than {@code name=value} + * @throws IllegalArgumentException if a parameter with that name has + * already been set + */ + public Builder withParameter(String name, String value) { + Objects.requireNonNull(name, "url parameter name cannot be null"); + if (parameters.containsKey(name)) { + throw new IllegalArgumentException("url parameter [" + name + "] has already been set to [" + parameters.get(name) + "]"); + } else { + parameters.put(name, value); + } + return this; + } + + /** + * Add query parameters using the provided map of key value pairs. + * + * @param paramSource a map of key value pairs where the key is the url parameter. + * @throws IllegalArgumentException if a parameter with that name has already been set. + */ + public Builder withParameters(Map paramSource) { + paramSource.forEach(this::withParameter); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions.Builder options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options.build(); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + return this; + } + + /** + * Set the body of the request. If not set or set to {@code null} then no + * body is sent with the request. + * + * @param entity the {@link BodyPublisher} to be set as the body of the request. + */ + public Builder withEntity(BodyPublisher entity) { + this.entity = entity; + return this; + } + + /** + * Set the body of the request to a string. If not set or set to + * {@code null} then no body is sent with the request. The + * {@code Content-Type} will be sent as {@code application/json}. + * If you need a different content type then use + * {@link #withEntity(BodyPublisher)}. + * + * @param entity JSON string to be set as the entity body of the request. + */ + public Builder withEntity(String entity) { + withEntity(entity == null ? BodyPublishers.noBody() : BodyPublishers.ofString(entity)); + return this; + } + + public Request build() { + return new Request(method, endpoint, parameters, entity, options); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java new file mode 100644 index 0000000000000..06899bbf86d5a --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLine.java @@ -0,0 +1,68 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.Serializable; +import java.net.URI; +import java.net.http.HttpClient.Version; +import java.net.http.HttpRequest; +import java.util.Objects; + +/** + * Request line (protocol, method, uri) + * Note: This is an experimental API. + */ +public record RequestLine(String method, String uri, Version protocolVersion) implements Serializable { + private static final long serialVersionUID = 2810581718468737193L; + + public RequestLine { + method = Objects.requireNonNull(method, "Method"); + } + + /** + * Create a new instance from the request + * @param request HTTP request + */ + public RequestLine(final HttpRequest request) { + this(Objects.requireNonNull(request, "Request").method(), buildUri(request.uri()), request.version().orElse(Version.HTTP_1_1)); + } + + private static String buildUri(URI uri) { + final String query = uri.getQuery(); + if (query != null && query.isBlank() == false) { + return uri.getPath() + "?" + query; + } else { + return uri.getPath(); + } + } + + /** + * Creates new request line instance + * @param method request HTTP method + * @param uri request uri + * @param version HTTP protocol + */ + public RequestLine(final String method, final URI uri, final Version version) { + this( + Objects.requireNonNull(method, "Method"), + Objects.requireNonNull(uri, "URI").getPath(), + version != null ? version : Version.HTTP_1_1 + ); + } + + /** + * Converts the request line to string + */ + @Override + public String toString() { + final StringBuilder buf = new StringBuilder(); + buf.append(this.method).append(" ").append(this.uri).append(" ").append(this.protocolVersion); + return buf.toString(); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java new file mode 100644 index 0000000000000..1120669b889ba --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestLogger.java @@ -0,0 +1,190 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import java.io.IOException; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Flow; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * Helper class that exposes static methods to unify the way requests are logged. + * Includes trace logging to log complete requests and responses in curl format. + * Useful for debugging, manually sending logged requests via curl and checking their responses. + * Trace logging is a feature that all the language clients provide. + */ +final class RequestLogger { + + private static final Log tracer = LogFactory.getLog("tracer"); + + private RequestLogger() {} + + /** + * Logs a streaming request that yielded a streaming response + */ + static void logStreamingResponse( + Log logger, + HttpRequest request, + HttpHost host, + HttpResponse>> httpResponse + ) { + logResponse(logger, request, host, httpResponse, List.of()); + } + + /** + * Logs a request that yielded a response + */ + static void logResponse(Log logger, HttpRequest request, HttpHost host, HttpResponse> httpResponse) { + logResponse(logger, request, host, httpResponse, httpResponse.body()); + } + + /** + * Logs a request that failed + */ + static void logFailedRequest(Log logger, Function request, Node node, Exception e) { + if (logger.isDebugEnabled()) { + final HttpRequest r = request.apply(node); + logger.debug("request [" + r.method() + " " + node.getHost() + r.uri() + "] failed", e); + } + if (tracer.isTraceEnabled()) { + String traceRequest; + try { + traceRequest = buildTraceRequest(request, node); + } catch (IOException e1) { + tracer.trace("error while reading request for trace purposes", e); + traceRequest = ""; + } + tracer.trace(traceRequest); + } + } + + static String buildWarningMessage(HttpRequest request, HttpHost host, List warnings) { + StringBuilder message = new StringBuilder("request [").append(request.method()) + .append(" ") + .append(host) + .append(request.uri()) + .append("] returned ") + .append(warnings.size()) + .append(" warnings: "); + for (int i = 0; i < warnings.size(); i++) { + if (i > 0) { + message.append(","); + } + message.append("[").append(warnings.get(i)).append("]"); + } + return message.toString(); + } + + /** + * Creates curl output for given request + */ + static String buildTraceRequest(Function request, Node node) throws IOException { + final HttpRequest r = request.apply(node); + return buildTraceRequest(r, node.getHost()); + } + + static String buildTraceRequest(HttpRequest request, HttpHost host) throws IOException { + String requestLine = "curl -iX " + request.method() + " '" + request.uri() + "'"; + final String body = BodyUtils.getBodyAsString(request).block(); + if (body != null) { + requestLine += " -d '"; + requestLine += body + "'"; + } + return requestLine; + } + + /** + * Creates curl output for given response + */ + static String buildTraceResponse(HttpResponse httpResponse, List body) throws IOException { + StringBuilder responseLine = new StringBuilder(); + responseLine.append("# ").append(new StatusLine(httpResponse)); + for (Map.Entry> header : httpResponse.headers().map().entrySet()) { + responseLine.append("\n# ") + .append(header.getKey()) + .append(": ") + .append(header.getValue().stream().collect(Collectors.joining(","))); + } + responseLine.append("\n#"); + + if (body != null && body.isEmpty() == false) { + for (ByteBuffer chunk : body) { + responseLine.append("\n# ").append(StandardCharsets.UTF_8.decode(chunk).toString()); + } + } + + return responseLine.toString(); + } + + /** + * Logs a request that yielded a response + */ + private static void logResponse(Log logger, HttpRequest request, HttpHost host, HttpResponse httpResponse, List body) { + if (logger.isDebugEnabled()) { + logger.debug("request [" + request.method() + " " + request.uri() + "] returned [" + new StatusLine(httpResponse) + "]"); + } + if (logger.isWarnEnabled()) { + List warnings = httpResponse.headers().allValues("Warning"); + if (warnings != null && warnings.size() > 0) { + logger.warn(buildWarningMessage(request, host, warnings)); + } + } + if (tracer.isTraceEnabled()) { + String requestLine; + try { + requestLine = buildTraceRequest(request, host); + } catch (IOException e) { + requestLine = ""; + tracer.trace("error while reading request for trace purposes", e); + } + String responseLine; + try { + responseLine = buildTraceResponse(httpResponse, body); + } catch (IOException e) { + responseLine = ""; + tracer.trace("error while reading response for trace purposes", e); + } + tracer.trace(requestLine + '\n' + responseLine); + } + } + +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java new file mode 100644 index 0000000000000..c814b257c6100 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RequestOptions.java @@ -0,0 +1,289 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * The portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * Note: This is an experimental API. + */ +public final class RequestOptions { + /** + * Default request options. + */ + public static final RequestOptions DEFAULT = new Builder( + Collections.emptyMap(), + Collections.emptyMap(), + null, + Duration.ofMillis(RestHttpClientBuilder.DEFAULT_RESPONSE_TIMEOUT_MILLIS) + ).build(); + + private final Map> headers; + private final Map parameters; + private final WarningsHandler warningsHandler; + private final Duration timeout; + + private RequestOptions(Builder builder) { + this.headers = Collections.unmodifiableMap(new HashMap<>(builder.headers)); + this.parameters = Collections.unmodifiableMap(new HashMap<>(builder.parameters)); + this.warningsHandler = builder.warningsHandler; + this.timeout = builder.timeout; + } + + /** + * Create a builder that contains these options but can be modified. + */ + public Builder toBuilder() { + return new Builder(headers, parameters, warningsHandler, timeout); + } + + /** + * Headers to attach to the request. + */ + public Map> getHeaders() { + return headers; + } + + /** + * Query parameters to attach to the request. Any parameters present here + * will override matching parameters in the {@link Request}, if they exist. + */ + public Map getParameters() { + return parameters; + } + + /** + * How this request should handle warnings. If null (the default) then + * this request will default to the behavior dictacted by + * {@link RestHttpClientBuilder#setStrictDeprecationMode}. + *

+ * This can be set to {@link WarningsHandler#PERMISSIVE} if the client + * should ignore all warnings which is the same behavior as setting + * strictDeprecationMode to true. It can be set to + * {@link WarningsHandler#STRICT} if the client should fail if there are + * any warnings which is the same behavior as settings + * strictDeprecationMode to false. + *

+ * It can also be set to a custom implementation of + * {@linkplain WarningsHandler} to permit only certain warnings or to + * fail the request if the warnings returned don't + * exactly match some set. + */ + public WarningsHandler getWarningsHandler() { + return warningsHandler; + } + + /** + * Gets the request timeout + * @return request timeout + */ + public Duration getTimeout() { + return timeout; + } + + /** + * Convert request options to string representation + */ + @Override + public String toString() { + StringBuilder b = new StringBuilder(); + b.append("RequestOptions{"); + boolean comma = false; + if (headers.size() > 0) { + comma = true; + b.append("headers="); + b.append(headers.entrySet().stream().map(e -> e.getKey() + "=" + e.getValue()).collect(Collectors.joining(","))); + } + if (timeout != null) { + if (comma) b.append(", "); + comma = true; + b.append("timeout=").append(timeout.toMillis()).append("ms"); + } + if (parameters.size() > 0) { + if (comma) b.append(", "); + comma = true; + b.append("parameters="); + b.append(parameters.entrySet().stream().map(e -> e.getKey() + "=" + e.getValue()).collect(Collectors.joining(","))); + } + if (warningsHandler != null) { + if (comma) b.append(", "); + comma = true; + b.append("warningsHandler=").append(warningsHandler); + } + return b.append('}').toString(); + } + + /** + * Compare two request options for equality + * @param obj request options instance to compare with + */ + @Override + public boolean equals(Object obj) { + if (obj == null || (obj.getClass() != getClass())) { + return false; + } + if (obj == this) { + return true; + } + + RequestOptions other = (RequestOptions) obj; + return headers.equals(other.headers) + && parameters.equals(other.parameters) + && Objects.equals(timeout, other.timeout) + && Objects.equals(warningsHandler, other.warningsHandler); + } + + /** + * Calculate the hash code of the request options + */ + @Override + public int hashCode() { + return Objects.hash(headers, parameters, warningsHandler); + } + + /** + * Builds {@link RequestOptions}. Get one by calling + * {@link RequestOptions#toBuilder} on {@link RequestOptions#DEFAULT} or + * any other {@linkplain RequestOptions}. + */ + public static class Builder { + private final Map> headers; + private final Map parameters; + private WarningsHandler warningsHandler; + private Duration timeout = Duration.ofMillis(RestHttpClientBuilder.DEFAULT_RESPONSE_TIMEOUT_MILLIS); + + private Builder( + Map> headers, + Map parameters, + WarningsHandler warningsHandler, + Duration timeout + ) { + this.headers = new HashMap<>(headers); + this.parameters = new HashMap<>(parameters); + this.warningsHandler = warningsHandler; + this.timeout = timeout; + } + + /** + * Build the {@linkplain RequestOptions}. + */ + public RequestOptions build() { + return new RequestOptions(this); + } + + /** + * Sets request timeout + * + * @param timeout timeout + */ + public void setTimeout(Duration timeout) { + Objects.requireNonNull(timeout, "timeout cannot be null"); + this.timeout = timeout; + } + + /** + * Add the provided headers to the request. + * + * @param headers headers to add + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addHeaders(Map> headers) { + Objects.requireNonNull(headers, "headers cannot be null"); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> addHeader(header.getKey(), v)); + } + return this; + } + + /** + * Add the provided header to the request. + * + * @param name the header name + * @param value the header value + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addHeader(String name, String value) { + Objects.requireNonNull(name, "header name cannot be null"); + Objects.requireNonNull(value, "header value cannot be null"); + this.headers.computeIfAbsent(name, v -> new ArrayList<>()).add(value); + return this; + } + + /** + * Add the provided query parameter to the request. Any parameters added here + * will override matching parameters in the {@link Request}, if they exist. + * + * @param name the query parameter name + * @param value the query parameter value + * @throws NullPointerException if {@code name} or {@code value} is null. + */ + public Builder addParameter(String name, String value) { + Objects.requireNonNull(name, "query parameter name cannot be null"); + Objects.requireNonNull(value, "query parameter value cannot be null"); + this.parameters.put(name, value); + return this; + } + + /** + * How this request should handle warnings. If null (the default) then + * this request will default to the behavior dictacted by + * {@link RestHttpClientBuilder#setStrictDeprecationMode}. + *

+ * This can be set to {@link WarningsHandler#PERMISSIVE} if the client + * should ignore all warnings which is the same behavior as setting + * strictDeprecationMode to true. It can be set to + * {@link WarningsHandler#STRICT} if the client should fail if there are + * any warnings which is the same behavior as settings + * strictDeprecationMode to false. + *

+ * It can also be set to a custom implementation of + * {@linkplain WarningsHandler} to permit only certain warnings or to + * fail the request if the warnings returned don't + * exactly match some set. + * + * @param warningsHandler the {@link WarningsHandler} to be used + */ + public void setWarningsHandler(WarningsHandler warningsHandler) { + this.warningsHandler = warningsHandler; + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java new file mode 100644 index 0000000000000..dccfb8d24a79b --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/Response.java @@ -0,0 +1,117 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.InputStream; +import java.net.http.HttpHeaders; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.Flow; + +/** + * Holds an OpenSearch response. It wraps the {@link HttpResponse} returned and associates it with + * its corresponding {@link RequestLine} and {@link HttpHost}. + * Note: This is an experimental API. + */ +public sealed interface Response permits CompressedResponse, NonCompressedResponse { + /** + * Create response from streaming conversation + * @param requestLine request line + * @param host host + * @param response underlying HTTP response + * @return new response instance + */ + static Response fromStreaming(RequestLine requestLine, HttpHost host, HttpResponse>> response) { + return new NonCompressedResponse(requestLine, host, response, List.of() /* streaming body could be very large */); + } + + /** + * Create response from non-streaming conversation + * @param requestLine request line + * @param host host + * @param response underlying HTTP response + * @return new response instance + */ + static Response from(RequestLine requestLine, HttpHost host, HttpResponse> response) { + final boolean compressed = response.headers() + .firstValue("Content-Encoding") + .filter("gzip"::equalsIgnoreCase) + .map(h -> true) + .orElse(false); + if (compressed == false) { + return new NonCompressedResponse(requestLine, host, response, response.body()); + } else { + return new CompressedResponse(requestLine, host, response, response.body()); + } + } + + /** + * Returns the request line that generated this response + */ + RequestLine requestLine(); + + /** + * Returns the node that returned this response + */ + HttpHost host(); + + /** + * Returns the status line of the current response + */ + default StatusLine statusLine() { + return new StatusLine(httpResponse()); + } + + /** + * Returns all the response headers + */ + default HttpHeaders headers() { + return httpResponse().headers(); + } + + /** + * Returns the response body available, null otherwise + * @see InputStream + */ + List entity(); + + /** + * Returns the value of the first header with a specified name of this message. + * If there is more than one matching header in the message the first element is returned. + * If there is no matching header in the message null is returned. + * + * @param name header name + */ + default String header(String name) { + return headers().firstValue(name).orElse(null); + } + + /** + * Returns a list of all warning headers returned in the response. + */ + default List warnings() { + return ResponseWarningsExtractor.getWarnings(httpResponse()); + } + + /** + * Returns true if there is at least one warning header returned in the + * response. + */ + default boolean hasWarnings() { + List warnings = headers().allValues("Warning"); + return warnings != null && warnings.size() > 0; + } + + /** + * Returns underlying HTTP response instance + * @return + */ + HttpResponse httpResponse(); +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java new file mode 100644 index 0000000000000..05d95bd9014ab --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseException.java @@ -0,0 +1,61 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Locale; + +/** + * Exception thrown when an opensearch node responds to a request with a status code that indicates an error. + * Holds the response that was returned. + */ +public final class ResponseException extends IOException { + private static final long serialVersionUID = 1L; + private final Response response; + + /** + * Creates a ResponseException containing the given {@code Response}. + * + * @param response The error response. + */ + public ResponseException(Response response) throws IOException { + super(buildMessage(response)); + this.response = response; + } + + static String buildMessage(Response response) throws IOException { + String message = String.format( + Locale.ROOT, + "method [%s], host [%s], URI [%s], status line [%s]", + response.requestLine().method(), + response.host(), + response.requestLine().uri(), + response.statusLine().toString() + ); + + if (response.hasWarnings()) { + message += "\n" + "Warnings: " + response.warnings(); + } + + List entity = response.entity(); + if (entity != null) { + message += "\n" + BodyUtils.getBodyAsString(response); + } + return message; + } + + /** + * Returns the {@link Response} that caused this exception to be thrown. + */ + public Response getResponse() { + return response; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java new file mode 100644 index 0000000000000..6775322ff5e71 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseListener.java @@ -0,0 +1,61 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +/** + * Listener to be provided when calling async performRequest methods provided by {@link RestHttpClient}. + * Those methods that do accept a listener will return immediately, execute asynchronously, and notify + * the listener whenever the request yielded a response, or failed with an exception. + * + *

+ * Note that it is not safe to call {@link RestHttpClient#close()} from either of these + * callbacks. + */ +public interface ResponseListener { + + /** + * Method invoked if the request yielded a successful response. + * + * @param response success response + */ + void onSuccess(Response response); + + /** + * Method invoked if the request failed. There are two main categories of failures: connection failures (usually + * {@link java.io.IOException}s, or responses that were treated as errors based on their error response code + * ({@link ResponseException}s). + * + * @param exception the failure exception + */ + void onFailure(Exception exception); +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java new file mode 100644 index 0000000000000..296538333fea3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/ResponseWarningsExtractor.java @@ -0,0 +1,96 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpResponse; +import java.util.ArrayList; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +final class ResponseWarningsExtractor { + + /** + * Optimized regular expression to test if a string matches the RFC 1123 date + * format (with quotes and leading space). Start/end of line characters and + * atomic groups are used to prevent backtracking. + */ + private static final Pattern WARNING_HEADER_DATE_PATTERN = Pattern.compile("^ " + // start of line, leading space + // quoted RFC 1123 date format + "\"" + // opening quote + "(?>Mon|Tue|Wed|Thu|Fri|Sat|Sun), " + // day of week, atomic group to prevent backtracking + "\\d{2} " + // 2-digit day + "(?>Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec) " + // month, atomic group to prevent backtracking + "\\d{4} " + // 4-digit year + "\\d{2}:\\d{2}:\\d{2} " + // (two-digit hour):(two-digit minute):(two-digit second) + "GMT" + // GMT + "\"$"); // closing quote (optional, since an older version can still send a warn-date), end of line + + /** + * Length of RFC 1123 format (with quotes and leading space), used in + * matchWarningHeaderPatternByPrefix(String). + */ + private static final int WARNING_HEADER_DATE_LENGTH = 0 + 1 + 1 + 3 + 1 + 1 + 2 + 1 + 3 + 1 + 4 + 1 + 2 + 1 + 2 + 1 + 2 + 1 + 3 + 1; + + private ResponseWarningsExtractor() {} + + /** + * Returns a list of all warning headers returned in the response. + * @param response HTTP response + */ + static List getWarnings(final HttpResponse response) { + List warnings = new ArrayList<>(); + for (String warning : response.headers().allValues("Warning")) { + if (matchWarningHeaderPatternByPrefix(warning)) { + warnings.add(extractWarningValueFromWarningHeader(warning)); + } else { + warnings.add(warning); + } + } + return warnings; + } + + /** + * Tests if a string matches the RFC 7234 specification for warning headers. + * This assumes that the warn code is always 299 and the warn agent is always + * OpenSearch. + * + * @param s the value of a warning header formatted according to RFC 7234 + * @return {@code true} if the input string matches the specification + */ + private static boolean matchWarningHeaderPatternByPrefix(final String s) { + return s.startsWith("299 OpenSearch-"); + } + + /** + * Refer to org.opensearch.common.logging.DeprecationLogger + */ + private static String extractWarningValueFromWarningHeader(final String s) { + String warningHeader = s; + + /* + * The following block tests for the existence of a RFC 1123 date in the warning header. If the date exists, it is removed for + * extractWarningValueFromWarningHeader(String) to work properly (as it does not handle dates). + */ + if (s.length() > WARNING_HEADER_DATE_LENGTH) { + final String possibleDateString = s.substring(s.length() - WARNING_HEADER_DATE_LENGTH); + final Matcher matcher = WARNING_HEADER_DATE_PATTERN.matcher(possibleDateString); + + if (matcher.matches()) { + warningHeader = warningHeader.substring(0, s.length() - WARNING_HEADER_DATE_LENGTH); + } + } + + final int firstQuote = warningHeader.indexOf('\"'); + final int lastQuote = warningHeader.length() - 1; + final String warningValue = warningHeader.substring(firstQuote + 1, lastQuote); + return warningValue; + } + +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java new file mode 100644 index 0000000000000..446b537d68826 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClient.java @@ -0,0 +1,1096 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import javax.net.ssl.SSLHandshakeException; + +import java.io.Closeable; +import java.io.IOException; +import java.net.ConnectException; +import java.net.SocketTimeoutException; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.net.http.HttpResponse; +import java.net.http.HttpResponse.BodyHandler; +import java.net.http.HttpResponse.BodyHandlers; +import java.net.http.HttpTimeoutException; +import java.nio.ByteBuffer; +import java.nio.channels.ClosedChannelException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Base64; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Flow; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; +import java.util.stream.Collectors; + +import org.reactivestreams.Publisher; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Collections.singletonList; + +/** + * Client that connects to an OpenSearch cluster through HTTP. + *

+ * Must be created using {@link RestHttpClientBuilder}, which allows to set all the different options or just rely on defaults. + * The hosts that are part of the cluster need to be provided at creation time, but can also be replaced later + * by calling {@link #setNodes(Collection)}. + *

+ * The method {@link #performRequest(Request)} allows to send a request to the cluster. When + * sending a request, a host gets selected out of the provided ones in a round-robin fashion. Failing hosts are marked dead and + * retried after a certain amount of time (minimum 1 minute, maximum 30 minutes), depending on how many times they previously + * failed (the more failures, the later they will be retried). In case of failures all of the alive nodes (or dead nodes that + * deserve a retry) are retried until one responds or none of them does, in which case an {@link IOException} will be thrown. + *

+ * Requests can be either synchronous or asynchronous. The asynchronous variants all end with {@code Async}. + *

+ * Requests can be traced by enabling trace logging for "tracer". The trace logger outputs requests and responses in curl format. + * + * Note: This is an experimental API. + */ +public class RestHttpClient implements Closeable { + + private static final Log logger = LogFactory.getLog(RestHttpClient.class); + + private final HttpClient client; + final Map> defaultHeaders; + private final String pathPrefix; + private final AtomicInteger lastNodeIndex = new AtomicInteger(0); + private final ConcurrentMap denylist = new ConcurrentHashMap<>(); + private final FailureListener failureListener; + private final NodeSelector nodeSelector; + private volatile List nodes; + private final WarningsHandler warningsHandler; + private final boolean compressionEnabled; + + RestHttpClient( + HttpClient client, + Map> defaultHeaders, + List nodes, + String pathPrefix, + FailureListener failureListener, + NodeSelector nodeSelector, + boolean strictDeprecationMode, + boolean compressionEnabled + ) { + this.client = client; + this.defaultHeaders = Collections.unmodifiableMap(defaultHeaders); + this.failureListener = failureListener; + this.pathPrefix = pathPrefix; + this.nodeSelector = nodeSelector; + this.warningsHandler = strictDeprecationMode ? WarningsHandler.STRICT : WarningsHandler.PERMISSIVE; + this.compressionEnabled = compressionEnabled; + setNodes(nodes); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the nodes that the client will send requests to. + * + * @param cloudId a valid elastic cloud cloudId that will route to a cluster. The cloudId is located in + * the user console https://cloud.elastic.co and will resemble a string like the following + * optionalHumanReadableName:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyRlbGFzdGljc2VhcmNoJGtpYmFuYQ== + */ + public static RestHttpClientBuilder builder(String cloudId) { + // there is an optional first portion of the cloudId that is a human readable string, but it is not used. + if (cloudId.contains(":")) { + if (cloudId.indexOf(":") == cloudId.length() - 1) { + throw new IllegalStateException("cloudId " + cloudId + " must begin with a human readable identifier followed by a colon"); + } + cloudId = cloudId.substring(cloudId.indexOf(":") + 1); + } + + String decoded = new String(Base64.getDecoder().decode(cloudId), UTF_8); + // once decoded the parts are separated by a $ character. + // they are respectively domain name and optional port, opensearch id, opensearch-dashboards id + String[] decodedParts = decoded.split("\\$"); + if (decodedParts.length != 3) { + throw new IllegalStateException("cloudId " + cloudId + " did not decode to a cluster identifier correctly"); + } + + // domain name and optional port + String[] domainAndMaybePort = decodedParts[0].split(":", 2); + String domain = domainAndMaybePort[0]; + int port; + + if (domainAndMaybePort.length == 2) { + try { + port = Integer.parseInt(domainAndMaybePort[1]); + } catch (NumberFormatException nfe) { + throw new IllegalStateException("cloudId " + cloudId + " does not contain a valid port number"); + } + } else { + port = 443; + } + + String url = decodedParts[1] + "." + domain; + return builder(new HttpHost("https", url, port)); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the hosts that the client will send requests to. + *

+ * Prefer this to {@link #builder(HttpHost...)} if you have metadata up front about the nodes. + * If you don't either one is fine. + * + * @param nodes The nodes that the client will send requests to. + */ + public static RestHttpClientBuilder builder(Node... nodes) { + return new RestHttpClientBuilder(nodes == null ? null : Arrays.asList(nodes)); + } + + /** + * Returns a new {@link RestHttpClientBuilder} to help with {@link RestHttpClient} creation. + * Creates a new builder instance and sets the nodes that the client will send requests to. + *

+ * You can use this if you do not have metadata up front about the nodes. If you do, prefer + * {@link #builder(Node...)}. + * @see Node#Node(HttpHost) + * + * @param hosts The hosts that the client will send requests to. + */ + public static RestHttpClientBuilder builder(HttpHost... hosts) { + if (hosts == null || hosts.length == 0) { + throw new IllegalArgumentException("hosts must not be null nor empty"); + } + List nodes = Arrays.stream(hosts).map(Node::new).collect(Collectors.toList()); + return new RestHttpClientBuilder(nodes); + } + + /** + * Replaces the nodes with which the client communicates. + * + * @param nodes the new nodes to communicate with. + */ + public synchronized void setNodes(Collection nodes) { + if (nodes == null || nodes.isEmpty()) { + throw new IllegalArgumentException("nodes must not be null or empty"); + } + + Map nodesByHost = new LinkedHashMap<>(); + for (Node node : nodes) { + Objects.requireNonNull(node, "node cannot be null"); + // TODO should we throw an IAE if we have two nodes with the same host? + nodesByHost.put(node.getHost(), node); + } + this.nodes = Collections.unmodifiableList(new ArrayList<>(nodesByHost.values())); + this.denylist.clear(); + } + + /** + * Get the list of nodes that the client knows about. The list is + * unmodifiable. + */ + public List getNodes() { + return nodes; + } + + /** + * check client running status + * @return client running status + */ + public boolean isRunning() { + return client.isTerminated() == false; + } + + /** + * Sends a streaming request to the OpenSearch cluster that the client points to and returns streaming response. This is an experimental API. + * @param request streaming request + * @return streaming response + * @throws IOException IOException + */ + public StreamingResponse streamRequest(StreamingRequest request) throws IOException { + final InternalStreamingRequest internalRequest = new InternalStreamingRequest(request); + return streamRequest(nextNodes(), internalRequest); + } + + /** + * Sends a request to the OpenSearch cluster that the client points to. + * Blocks until the request is completed and returns its response or fails + * by throwing an exception. Selects a host out of the provided ones in a + * round-robin fashion. Failing hosts are marked dead and retried after a + * certain amount of time (minimum 1 minute, maximum 30 minutes), depending + * on how many times they previously failed (the more failures, the later + * they will be retried). In case of failures all of the alive nodes (or + * dead nodes that deserve a retry) are retried until one responds or none + * of them does, in which case an {@link IOException} will be thrown. + *

+ * This method works by performing an asynchronous call and waiting + * for the result. If the asynchronous call throws an exception we wrap + * it and rethrow it so that the stack trace attached to the exception + * contains the call site. While we attempt to preserve the original + * exception this isn't always possible and likely haven't covered all of + * the cases. You can get the original exception from + * {@link Exception#getCause()}. + * + * @param request the request to perform + * @return the response returned by OpenSearch + * @throws IOException in case of a problem or the connection was aborted + * @throws ResponseException in case OpenSearch responded with a status code that indicated an error + */ + public Response performRequest(Request request) throws IOException { + InternalRequest internalRequest = new InternalRequest(request); + return performRequest(nextNodes(), internalRequest, null); + } + + private Response performRequest(final Iterator nodes, final InternalRequest request, Exception previousException) + throws IOException { + Node node = nodes.next(); + RequestContext> context = request.createContextForNextAttempt(node); + HttpResponse> httpResponse; + try { + httpResponse = client.send(context.requestProducer(), context.asyncResponseConsumer()); + } catch (Exception e) { + RequestLogger.logFailedRequest(logger, request.httpRequest, context.node(), e); + onFailure(context.node()); + Exception cause = extractAndWrapCause(e); + addSuppressedException(previousException, cause); + if (nodes.hasNext()) { + return performRequest(nodes, request, cause); + } + if (cause instanceof IOException) { + throw (IOException) cause; + } + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + throw new IllegalStateException("unexpected exception type: must be either RuntimeException or IOException", cause); + } + ResponseOrResponseException responseOrResponseException = convertResponse(request, context.node(), httpResponse); + if (responseOrResponseException.responseException == null) { + return responseOrResponseException.response; + } + addSuppressedException(previousException, responseOrResponseException.responseException); + if (nodes.hasNext()) { + return performRequest(nodes, request, responseOrResponseException.responseException); + } + throw responseOrResponseException.responseException; + } + + private Publisher>>> streamRequest( + final Node node, + final Iterator nodes, + final InternalStreamingRequest request + ) throws IOException { + return request.cancellable.callIfNotCancelled(() -> { + final RequestContext>> context = request.createContextForNextAttempt(node); + final CompletableFuture>>> future = client.sendAsync( + context.requestProducer(), + context.asyncResponseConsumer() + ); + request.setCancellable(future); + + final Mono>>> publisher = Mono.fromCompletionStage(future).flatMap(response -> { + try { + final ResponseOrResponseException responseOrResponseException = convertResponse(request, node, response); + if (responseOrResponseException.responseException == null) { + return Mono.just(response); + } else { + if (nodes.hasNext()) { + return Mono.from(streamRequest(nodes.next(), nodes, request)); + } else { + return Mono.error(responseOrResponseException.responseException); + } + } + } catch (final Exception ex) { + return Mono.error(ex); + } + }); + + return publisher; + }); + } + + private StreamingResponse streamRequest(final Iterator nodes, final InternalStreamingRequest request) throws IOException { + return request.cancellable.callIfNotCancelled(() -> { + final Node node = nodes.next(); + return new StreamingResponse(new RequestLine(request.httpRequest.apply(node)), streamRequest(node, nodes, request)); + }); + } + + private ResponseOrResponseException convertResponse(InternalRequest request, Node node, HttpResponse> httpResponse) + throws IOException { + + final HttpRequest httpRequest = request.httpRequest.apply(node); + RequestLogger.logResponse(logger, httpRequest, node.getHost(), httpResponse); + int statusCode = httpResponse.statusCode(); + + Response response = Response.from(new RequestLine(httpRequest), node.getHost(), httpResponse); + if (isSuccessfulResponse(statusCode) || request.ignoreErrorCodes.contains(response.statusLine().statusCode())) { + onResponse(node); + if (request.warningsHandler.warningsShouldFailRequest(response.warnings())) { + throw new WarningFailureException(response); + } + return new ResponseOrResponseException(response); + } + ResponseException responseException = new ResponseException(response); + if (isRetryStatus(statusCode)) { + // mark host dead and retry against next one + onFailure(node); + return new ResponseOrResponseException(responseException); + } + // mark host alive and don't retry, as the error should be a request problem + onResponse(node); + throw responseException; + } + + private ResponseOrResponseException convertResponse( + InternalStreamingRequest request, + Node node, + HttpResponse>> httpResponse + ) throws IOException { + + // Streaming Response could accumulate a lot of data so we may not be able to fully consume it. + final HttpRequest httpRequest = request.httpRequest.apply(node); + final Response response = Response.fromStreaming(new RequestLine(httpRequest), node.getHost(), httpResponse); + + RequestLogger.logStreamingResponse(logger, request.httpRequest.apply(node), node.getHost(), httpResponse); + int statusCode = httpResponse.statusCode(); + + if (isSuccessfulResponse(statusCode) || request.ignoreErrorCodes.contains(response.statusLine().statusCode())) { + onResponse(node); + if (request.warningsHandler.warningsShouldFailRequest(response.warnings())) { + throw new WarningFailureException(response); + } + return new ResponseOrResponseException(response); + } + ResponseException responseException = new ResponseException(response); + if (isRetryStatus(statusCode)) { + // mark host dead and retry against next one + onFailure(node); + return new ResponseOrResponseException(responseException); + } + // mark host alive and don't retry, as the error should be a request problem + onResponse(node); + throw responseException; + } + + /** + * Sends a request to the OpenSearch cluster that the client points to. + * The request is executed asynchronously and the provided + * {@link ResponseListener} gets notified upon request completion or + * failure. Selects a host out of the provided ones in a round-robin + * fashion. Failing hosts are marked dead and retried after a certain + * amount of time (minimum 1 minute, maximum 30 minutes), depending on how + * many times they previously failed (the more failures, the later they + * will be retried). In case of failures all of the alive nodes (or dead + * nodes that deserve a retry) are retried until one responds or none of + * them does, in which case an {@link IOException} will be thrown. + * + * @param request the request to perform + * @param responseListener the {@link ResponseListener} to notify when the + * request is completed or fails + */ + public Cancellable performRequestAsync(Request request, ResponseListener responseListener) { + try { + FailureTrackingResponseListener failureTrackingResponseListener = new FailureTrackingResponseListener(responseListener); + InternalRequest internalRequest = new InternalRequest(request); + performRequestAsync(nextNodes(), internalRequest, failureTrackingResponseListener); + return internalRequest.cancellable; + } catch (Exception e) { + responseListener.onFailure(e); + return Cancellable.NO_OP; + } + } + + private void performRequestAsync( + final Iterator nodes, + final InternalRequest request, + final FailureTrackingResponseListener listener + ) { + request.cancellable.runIfNotCancelled(() -> { + final RequestContext> context = request.createContextForNextAttempt(nodes.next()); + CompletableFuture>> future = client.sendAsync( + context.requestProducer(), + context.asyncResponseConsumer() + ); + + request.setCancellable(future); + future.whenComplete((httpResponse, throwable) -> { + if (httpResponse != null) { + try { + ResponseOrResponseException responseOrResponseException = convertResponse(request, context.node(), httpResponse); + if (responseOrResponseException.responseException == null) { + listener.onSuccess(responseOrResponseException.response); + } else { + if (nodes.hasNext()) { + listener.trackFailure(responseOrResponseException.responseException); + performRequestAsync(nodes, request, listener); + } else { + listener.onDefinitiveFailure(responseOrResponseException.responseException); + } + } + } catch (Exception e) { + listener.onDefinitiveFailure(e); + } + } else if (throwable instanceof Exception failure) { + if (failure instanceof CancellationException) { + listener.onDefinitiveFailure(Cancellable.newCancellationException()); + } else { + Exception cause = failure; + if (failure instanceof CompletionException && failure.getCause() instanceof Exception ce) { + cause = ce; + } + try { + RequestLogger.logFailedRequest(logger, request.httpRequest, context.node(), failure); + onFailure(context.node()); + if (nodes.hasNext()) { + listener.trackFailure(failure); + performRequestAsync(nodes, request, listener); + } else { + listener.onDefinitiveFailure(cause); + } + } catch (Exception e) { + listener.onDefinitiveFailure(e); + } + } + } + }); + }); + + } + + /** + * Returns a non-empty {@link Iterator} of nodes to be used for a request + * that match the {@link NodeSelector}. + *

+ * If there are no living nodes that match the {@link NodeSelector} + * this will return the dead node that matches the {@link NodeSelector} + * that is closest to being revived. + * @throws IOException if no nodes are available + */ + private Iterator nextNodes() throws IOException { + List nodes = this.nodes; + Iterable hosts = selectNodes(nodes, denylist, lastNodeIndex, nodeSelector); + return hosts.iterator(); + } + + /** + * Select nodes to try and sorts them so that the first one will be tried initially, then the following ones + * if the previous attempt failed and so on. Package private for testing. + */ + static Iterable selectNodes( + List nodes, + Map denylist, + AtomicInteger lastNodeIndex, + NodeSelector nodeSelector + ) throws IOException { + /* + * Sort the nodes into living and dead lists. + */ + List livingNodes = new ArrayList<>(Math.max(0, nodes.size() - denylist.size())); + List deadNodes = new ArrayList<>(denylist.size()); + for (Node node : nodes) { + DeadHostState deadness = denylist.get(node.getHost()); + if (deadness == null || deadness.shallBeRetried()) { + livingNodes.add(node); + } else { + deadNodes.add(new DeadNode(node, deadness)); + } + } + + if (false == livingNodes.isEmpty()) { + /* + * Normal state: there is at least one living node. If the + * selector is ok with any over the living nodes then use them + * for the request. + */ + List selectedLivingNodes = new ArrayList<>(livingNodes); + nodeSelector.select(selectedLivingNodes); + if (false == selectedLivingNodes.isEmpty()) { + /* + * Rotate the list using a global counter as the distance so subsequent + * requests will try the nodes in a different order. + */ + Collections.rotate(selectedLivingNodes, lastNodeIndex.getAndIncrement()); + return selectedLivingNodes; + } + } + + /* + * Last resort: there are no good nodes to use, either because + * the selector rejected all the living nodes or because there aren't + * any living ones. Either way, we want to revive a single dead node + * that the NodeSelectors are OK with. We do this by passing the dead + * nodes through the NodeSelector so it can have its say in which nodes + * are ok. If the selector is ok with any of the nodes then we will take + * the one in the list that has the lowest revival time and try it. + */ + if (false == deadNodes.isEmpty()) { + final List selectedDeadNodes = new ArrayList<>(deadNodes); + /* + * We'd like NodeSelectors to remove items directly from deadNodes + * so we can find the minimum after it is filtered without having + * to compare many things. This saves us a sort on the unfiltered + * list. + */ + nodeSelector.select(() -> new DeadNodeIteratorAdapter(selectedDeadNodes.iterator())); + if (false == selectedDeadNodes.isEmpty()) { + return singletonList(Collections.min(selectedDeadNodes).node); + } + } + throw new IOException( + "NodeSelector [" + nodeSelector + "] rejected all nodes, " + "living " + livingNodes + " and dead " + deadNodes + ); + } + + /** + * Called after each successful request call. + * Receives as an argument the host that was used for the successful request. + */ + private void onResponse(Node node) { + DeadHostState removedHost = this.denylist.remove(node.getHost()); + if (logger.isDebugEnabled() && removedHost != null) { + logger.debug("removed [" + node + "] from denylist"); + } + } + + /** + * Called after each failed attempt. + * Receives as an argument the host that was used for the failed attempt. + */ + private void onFailure(Node node) { + while (true) { + DeadHostState previousDeadHostState = denylist.putIfAbsent( + node.getHost(), + new DeadHostState(DeadHostState.DEFAULT_TIME_SUPPLIER) + ); + if (previousDeadHostState == null) { + if (logger.isDebugEnabled()) { + logger.debug("added [" + node + "] to denylist"); + } + break; + } + if (denylist.replace(node.getHost(), previousDeadHostState, new DeadHostState(previousDeadHostState))) { + if (logger.isDebugEnabled()) { + logger.debug("updated [" + node + "] already in denylist"); + } + break; + } + } + failureListener.onFailure(node); + } + + /** + * Close the underlying {@link HttpClient} instance + */ + @Override + public void close() throws IOException { + client.shutdownNow(); + client.close(); + } + + private static boolean isSuccessfulResponse(int statusCode) { + return statusCode < 300; + } + + private static boolean isRetryStatus(int statusCode) { + switch (statusCode) { + case 502: + case 503: + case 504: + return true; + } + return false; + } + + private static void addSuppressedException(Exception suppressedException, Exception currentException) { + if (suppressedException != null) { + currentException.addSuppressed(suppressedException); + } + } + + private HttpRequest.Builder createHttpRequest( + Node node, + String method, + URI uri, + BodyPublisher body, + Duration timeout, + boolean compressed + ) { + return HttpRequest.newBuilder() + .uri(URI.create(node.getHost().toString()).resolve(uri)) + .timeout(timeout) + .method( + method, + (body == null) ? BodyPublishers.noBody() + : compressed == false ? body + : BodyPublishers.fromPublisher( + JdkFlowAdapter.publisherToFlowPublisher( + JdkFlowAdapter.flowPublisherToFlux(body).buffer().map(BodyUtils::compress).flatMap(Flux::fromIterable) + ) + ) + ); + } + + private HttpRequest.Builder createStreamingHttpRequest( + Node node, + String method, + URI uri, + Publisher body, + Duration timeout, + boolean compressed + ) { + return HttpRequest.newBuilder() + .uri(URI.create(node.getHost().toString()).resolve(uri)) + .timeout(timeout) + .method( + method, + (body == null) ? BodyPublishers.noBody() + : compressed == false ? BodyPublishers.fromPublisher(JdkFlowAdapter.publisherToFlowPublisher(body)) + : BodyPublishers.fromPublisher(JdkFlowAdapter.publisherToFlowPublisher(Flux.from(body).map(BodyUtils::compress))) + ); + } + + static URI buildUri(String pathPrefix, String path, Map params) { + Objects.requireNonNull(path, "path must not be null"); + try { + String fullPath; + if (pathPrefix != null && pathPrefix.isEmpty() == false) { + if (pathPrefix.endsWith("/") && path.startsWith("/")) { + fullPath = pathPrefix.substring(0, pathPrefix.length() - 1) + path; + } else if (pathPrefix.endsWith("/") || path.startsWith("/")) { + fullPath = pathPrefix + path; + } else { + fullPath = pathPrefix + "/" + path; + } + } else { + fullPath = path; + } + + final String additionalQuery = params.entrySet().stream().map(e -> { + if (e.getValue() == null) { + return e.getKey(); + } else { + return e.getKey() + "=" + URLEncoder.encode(e.getValue(), StandardCharsets.UTF_8); + } + }).collect(Collectors.joining("&")); + final URI uri = URI.create(fullPath); + + String newQuery = uri.getQuery(); + if (newQuery == null) { + newQuery = additionalQuery; + } else { + newQuery += "&" + additionalQuery; + } + + return new URI(uri.getScheme(), uri.getAuthority(), uri.getPath(), newQuery, uri.getFragment()); + } catch (URISyntaxException e) { + throw new IllegalArgumentException(e.getMessage(), e); + } + } + + /** + * Listener used in any async call to wrap the provided user listener (or SyncResponseListener in sync calls). + * Allows to track potential failures coming from the different retry attempts and returning to the original listener + * only when we got a response (successful or not to be retried) or there are no hosts to retry against. + */ + static class FailureTrackingResponseListener { + private final ResponseListener responseListener; + private volatile Exception exception; + + FailureTrackingResponseListener(ResponseListener responseListener) { + this.responseListener = responseListener; + } + + /** + * Notifies the caller of a response through the wrapped listener + */ + void onSuccess(Response response) { + responseListener.onSuccess(response); + } + + /** + * Tracks one last definitive failure and returns to the caller by notifying the wrapped listener + */ + void onDefinitiveFailure(Exception exception) { + trackFailure(exception); + responseListener.onFailure(this.exception); + } + + /** + * Tracks an exception, which caused a retry hence we should not return yet to the caller + */ + void trackFailure(Exception exception) { + addSuppressedException(this.exception, exception); + this.exception = exception; + } + } + + /** + * Listener that allows to be notified whenever a failure happens. Useful when sniffing is enabled, so that we can sniff on failure. + * The default implementation is a no-op. + */ + public static class FailureListener { + /** + * Create a {@link FailureListener} instance. + */ + public FailureListener() {} + + /** + * Notifies that the node provided as argument has just failed. + * + * @param node The node which has failed. + */ + public void onFailure(Node node) {} + } + + /** + * Contains a reference to a denylisted node and the time until it is + * revived. We use this so we can do a single pass over the denylist. + */ + private static class DeadNode implements Comparable { + final Node node; + final DeadHostState deadness; + + DeadNode(Node node, DeadHostState deadness) { + this.node = node; + this.deadness = deadness; + } + + @Override + public String toString() { + return node.toString(); + } + + @Override + public int compareTo(DeadNode rhs) { + return deadness.compareTo(rhs.deadness); + } + } + + /** + * Adapts an Iterator<DeadNodeAndRevival> into an + * Iterator<Node>. + */ + private static class DeadNodeIteratorAdapter implements Iterator { + private final Iterator itr; + + private DeadNodeIteratorAdapter(Iterator itr) { + this.itr = itr; + } + + @Override + public boolean hasNext() { + return itr.hasNext(); + } + + @Override + public Node next() { + return itr.next().node; + } + + @Override + public void remove() { + itr.remove(); + } + } + + private class InternalStreamingRequest { + private final Set ignoreErrorCodes; + private final Function httpRequest; + private final WarningsHandler warningsHandler; + private volatile Cancellable cancellable = Cancellable.fromFuture(new CompletableFuture<>()); + + InternalStreamingRequest(StreamingRequest request) { + Map params = new HashMap<>(request.parameters()); + // ignore is a special parameter supported by the clients, shouldn't be sent to es + String ignoreString = params.remove("ignore"); + this.ignoreErrorCodes = getIgnoreErrorCodes(ignoreString, request.method()); + + this.httpRequest = node -> { + URI uri = buildUri(pathPrefix, request.endpoint(), params); + HttpRequest.Builder builder = createStreamingHttpRequest( + node, + request.method(), + uri, + request.body(), + request.options().getTimeout(), + compressionEnabled + ); + setHeaders(builder, request.options().getHeaders()); + return builder.build(); + }; + + this.warningsHandler = request.options().getWarningsHandler() == null + ? RestHttpClient.this.warningsHandler + : request.options().getWarningsHandler(); + } + + private void setHeaders(HttpRequest.Builder httpRequest, Map> requestHeaders) { + // request headers override default headers, so we don't add default headers if they exist as request headers + final Set requestNames = new HashSet<>(requestHeaders.size()); + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> httpRequest.header(requestHeader.getKey(), v)); + requestNames.add(requestHeader.getKey()); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + if (requestNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> httpRequest.header(defaultHeader.getKey(), v)); + } + } + if (compressionEnabled) { + httpRequest.header("Content-Encoding", "gzip"); + httpRequest.header("Accept-Encoding", "gzip"); + } + } + + private void setCancellable(Future f) { + cancellable = Cancellable.fromFuture(f); + } + + RequestContext>> createContextForNextAttempt(Node node) { + return new ReactiveRequestContext(this, node); + } + } + + private class InternalRequest { + private final Set ignoreErrorCodes; + private final Function httpRequest; + private final WarningsHandler warningsHandler; + private volatile Cancellable cancellable = Cancellable.fromFuture(new CompletableFuture<>()); + + InternalRequest(Request request) { + Map params = new HashMap<>(request.parameters()); + // ignore is a special parameter supported by the clients, shouldn't be sent to es + String ignoreString = params.remove("ignore"); + this.ignoreErrorCodes = getIgnoreErrorCodes(ignoreString, request.method()); + this.httpRequest = node -> { + URI uri = buildUri(pathPrefix, request.endpoint(), params); + final HttpRequest.Builder builder = createHttpRequest( + node, + request.method(), + uri, + request.entity(), + request.options().getTimeout(), + compressionEnabled + ); + setHeaders(builder, request.options().getHeaders()); + return builder.build(); + }; + this.warningsHandler = request.options().getWarningsHandler() == null + ? RestHttpClient.this.warningsHandler + : request.options().getWarningsHandler(); + } + + private void setCancellable(Future f) { + cancellable = Cancellable.fromFuture(f); + } + + private void setHeaders(HttpRequest.Builder httpRequest, Map> requestHeaders) { + // request headers override default headers, so we don't add default headers if they exist as request headers + final Set requestNames = new HashSet<>(requestHeaders.size()); + + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> httpRequest.header(requestHeader.getKey(), v)); + requestNames.add(requestHeader.getKey()); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + if (requestNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> httpRequest.header(defaultHeader.getKey(), v)); + } + } + if (compressionEnabled) { + httpRequest.header("Content-Encoding", "gzip"); + httpRequest.header("Accept-Encoding", "gzip"); + } + } + + RequestContext> createContextForNextAttempt(Node node) { + return new AsyncRequestContext(this, node); + } + } + + private interface RequestContext { + Node node(); + + HttpRequest requestProducer(); + + BodyHandler asyncResponseConsumer(); + } + + private static class ReactiveRequestContext implements RequestContext>> { + private final Node node; + private final HttpRequest requestProducer; + private final BodyHandler>> asyncResponseConsumer; + + ReactiveRequestContext(InternalStreamingRequest request, Node node) { + this.node = node; + // we stream the request body if the entity allows for it + this.requestProducer = request.httpRequest.apply(node); + this.asyncResponseConsumer = BodyHandlers.ofPublisher(); + } + + @Override + public BodyHandler>> asyncResponseConsumer() { + return asyncResponseConsumer; + } + + @Override + public Node node() { + return node; + } + + @Override + public HttpRequest requestProducer() { + return requestProducer; + } + } + + private static class AsyncRequestContext implements RequestContext> { + private final Node node; + private final HttpRequest requestProducer; + private final BodyHandler> asyncResponseConsumer; + + AsyncRequestContext(InternalRequest request, Node node) { + this.node = node; + this.requestProducer = request.httpRequest.apply(node); + this.asyncResponseConsumer = BodyHandlers.fromSubscriber(new AsyncResponseProducer(), AsyncResponseProducer::getResult); + } + + @Override + public BodyHandler> asyncResponseConsumer() { + return asyncResponseConsumer; + } + + @Override + public Node node() { + return node; + } + + @Override + public HttpRequest requestProducer() { + return requestProducer; + } + } + + private static Set getIgnoreErrorCodes(String ignoreString, String requestMethod) { + Set ignoreErrorCodes; + if (ignoreString == null) { + if ("HEAD".equalsIgnoreCase(requestMethod)) { + // 404 never causes error if returned for a HEAD request + ignoreErrorCodes = Collections.singleton(404); + } else { + ignoreErrorCodes = Collections.emptySet(); + } + } else { + String[] ignoresArray = ignoreString.split(","); + ignoreErrorCodes = new HashSet<>(); + if ("HEAD".equalsIgnoreCase(requestMethod)) { + // 404 never causes error if returned for a HEAD request + ignoreErrorCodes.add(404); + } + for (String ignoreCode : ignoresArray) { + try { + ignoreErrorCodes.add(Integer.valueOf(ignoreCode)); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("ignore value should be a number, found [" + ignoreString + "] instead", e); + } + } + } + return ignoreErrorCodes; + } + + private static class ResponseOrResponseException { + private final Response response; + private final ResponseException responseException; + + ResponseOrResponseException(Response response) { + this.response = Objects.requireNonNull(response); + this.responseException = null; + } + + ResponseOrResponseException(ResponseException responseException) { + this.responseException = Objects.requireNonNull(responseException); + this.response = null; + } + } + + /** + * Wrap the exception so the caller's signature shows up in the stack trace, taking care to copy the original type and message + * where possible so async and sync code don't have to check different exceptions. + */ + private static Exception extractAndWrapCause(Exception exception) { + if (exception instanceof InterruptedException) { + throw new RuntimeException("thread waiting for the response was interrupted", exception); + } + if (exception instanceof ExecutionException || exception instanceof CompletionException) { + Throwable t = exception.getCause() == null ? exception : exception.getCause(); + if (t instanceof Error) { + throw (Error) t; + } + exception = (Exception) t; + } + if (exception instanceof HttpTimeoutException) { + HttpTimeoutException e = new HttpTimeoutException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof ClosedChannelException) { + ClosedChannelException e = new ClosedChannelException(); + e.initCause(exception); + return e; + } + if (exception instanceof SocketTimeoutException) { + SocketTimeoutException e = new SocketTimeoutException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof SSLHandshakeException) { + SSLHandshakeException e = new SSLHandshakeException( + exception.getMessage() + "\nSee https://opensearch.org/docs/latest/clients/java-rest-high-level/ for troubleshooting." + ); + e.initCause(exception); + return e; + } + if (exception instanceof ConnectException) { + ConnectException e = new ConnectException(exception.getMessage()); + e.initCause(exception); + return e; + } + if (exception instanceof IOException) { + return new IOException(exception.getMessage(), exception); + } + if (exception instanceof RuntimeException) { + return new RuntimeException(exception.getMessage(), exception); + } + return new RuntimeException("error while performing request", exception); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java new file mode 100644 index 0000000000000..367e2efc05da4 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/RestHttpClientBuilder.java @@ -0,0 +1,248 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import javax.net.ssl.SSLContext; + +import java.net.Authenticator; +import java.net.http.HttpClient; +import java.security.AccessController; +import java.security.NoSuchAlgorithmException; +import java.security.PrivilegedAction; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * Helps creating a new {@link RestHttpClient}. Allows to set the most common http client configuration options when internally + * creating the underlying {@link HttpClient}. Also allows to provide an externally created + * {@link HttpClient} in case additional customization is needed. + * + * Note: This is an experimental API. + */ +public final class RestHttpClientBuilder { + /** + * The default connection timeout in milliseconds. + */ + public static final int DEFAULT_CONNECT_TIMEOUT_MILLIS = 1000; + + /** + * The default response timeout in milliseconds. + */ + public static final int DEFAULT_RESPONSE_TIMEOUT_MILLIS = 30000; + + private final List nodes; + private Map> defaultHeaders = Map.of(); + private RestHttpClient.FailureListener failureListener; + private HttpClientConfigCallback httpClientConfigCallback; + private String pathPrefix; + private NodeSelector nodeSelector = NodeSelector.ANY; + private boolean strictDeprecationMode = false; + private boolean compressionEnabled = false; + + /** + * Creates a new builder instance and sets the hosts that the client will send requests to. + * + * @throws IllegalArgumentException if {@code nodes} is {@code null} or empty. + */ + RestHttpClientBuilder(List nodes) { + if (nodes == null || nodes.isEmpty()) { + throw new IllegalArgumentException("nodes must not be null or empty"); + } + for (Node node : nodes) { + if (node == null) { + throw new IllegalArgumentException("node cannot be null"); + } + } + this.nodes = nodes; + } + + /** + * Sets the default request headers, which will be sent along with each request. + *

+ * Request-time headers will always overwrite any default headers. + * + * @param defaultHeaders array of default header + * @throws NullPointerException if {@code defaultHeaders} or any header is {@code null}. + */ + public RestHttpClientBuilder setDefaultHeaders(Map> defaultHeaders) { + Objects.requireNonNull(defaultHeaders, "defaultHeaders must not be null"); + this.defaultHeaders = Collections.unmodifiableMap(defaultHeaders); + return this; + } + + /** + * Sets the {@link RestHttpClient.FailureListener} to be notified for each request failure + * + * @param failureListener the {@link RestHttpClient.FailureListener} for each failure + * @throws NullPointerException if {@code failureListener} is {@code null}. + */ + public RestHttpClientBuilder setFailureListener(RestHttpClient.FailureListener failureListener) { + Objects.requireNonNull(failureListener, "failureListener must not be null"); + this.failureListener = failureListener; + return this; + } + + /** + * Sets the {@link HttpClientConfigCallback} to be used to customize http client configuration + * + * @param httpClientConfigCallback the {@link HttpClientConfigCallback} to be used + * @throws NullPointerException if {@code httpClientConfigCallback} is {@code null}. + */ + public RestHttpClientBuilder setHttpClientConfigCallback(HttpClientConfigCallback httpClientConfigCallback) { + Objects.requireNonNull(httpClientConfigCallback, "httpClientConfigCallback must not be null"); + this.httpClientConfigCallback = httpClientConfigCallback; + return this; + } + + /** + * Sets the path's prefix for every request used by the http client. + *

+ * For example, if this is set to "/my/path", then any client request will become "/my/path/" + endpoint. + *

+ * In essence, every request's {@code endpoint} is prefixed by this {@code pathPrefix}. The path prefix is useful for when + * OpenSearch is behind a proxy that provides a base path or a proxy that requires all paths to start with '/'; + * it is not intended for other purposes and it should not be supplied in other scenarios. + * + * @param pathPrefix the path prefix for every request. + * @throws NullPointerException if {@code pathPrefix} is {@code null}. + * @throws IllegalArgumentException if {@code pathPrefix} is empty, or ends with more than one '/'. + */ + public RestHttpClientBuilder setPathPrefix(String pathPrefix) { + this.pathPrefix = cleanPathPrefix(pathPrefix); + return this; + } + + /** + * Cleans up the given path prefix to ensure that looks like "/base/path". + * + * @param pathPrefix the path prefix to be cleaned up. + * @return the cleaned up path prefix. + * @throws NullPointerException if {@code pathPrefix} is {@code null}. + * @throws IllegalArgumentException if {@code pathPrefix} is empty, or ends with more than one '/'. + */ + public static String cleanPathPrefix(String pathPrefix) { + Objects.requireNonNull(pathPrefix, "pathPrefix must not be null"); + + if (pathPrefix.isEmpty()) { + throw new IllegalArgumentException("pathPrefix must not be empty"); + } + + String cleanPathPrefix = pathPrefix; + if (cleanPathPrefix.startsWith("/") == false) { + cleanPathPrefix = "/" + cleanPathPrefix; + } + + // best effort to ensure that it looks like "/base/path" rather than "/base/path/" + if (cleanPathPrefix.endsWith("/") && cleanPathPrefix.length() > 1) { + cleanPathPrefix = cleanPathPrefix.substring(0, cleanPathPrefix.length() - 1); + + if (cleanPathPrefix.endsWith("/")) { + throw new IllegalArgumentException("pathPrefix is malformed. too many trailing slashes: [" + pathPrefix + "]"); + } + } + return cleanPathPrefix; + } + + /** + * Sets the {@link NodeSelector} to be used for all requests. + * + * @param nodeSelector the {@link NodeSelector} to be used + * @throws NullPointerException if the provided nodeSelector is null + */ + public RestHttpClientBuilder setNodeSelector(NodeSelector nodeSelector) { + Objects.requireNonNull(nodeSelector, "nodeSelector must not be null"); + this.nodeSelector = nodeSelector; + return this; + } + + /** + * Whether the REST client should return any response containing at least + * one warning header as a failure. + * + * @param strictDeprecationMode flag for enabling strict deprecation mode + */ + public RestHttpClientBuilder setStrictDeprecationMode(boolean strictDeprecationMode) { + this.strictDeprecationMode = strictDeprecationMode; + return this; + } + + /** + * Whether the REST client should compress requests using gzip content encoding and add the "Accept-Encoding: gzip" + * header to receive compressed responses. + * + * @param compressionEnabled flag for enabling compression + */ + public RestHttpClientBuilder setCompressionEnabled(boolean compressionEnabled) { + this.compressionEnabled = compressionEnabled; + return this; + } + + /** + * Creates a new {@link RestHttpClient} based on the provided configuration. + */ + public RestHttpClient build() { + if (failureListener == null) { + failureListener = new RestHttpClient.FailureListener(); + } + @SuppressWarnings("removal") + HttpClient httpClient = AccessController.doPrivileged((PrivilegedAction) this::createHttpClient); + + return new RestHttpClient( + httpClient, + defaultHeaders, + nodes, + pathPrefix, + failureListener, + nodeSelector, + strictDeprecationMode, + compressionEnabled + ); + } + + @SuppressWarnings("removal") + private HttpClient createHttpClient() { + try { + HttpClient.Builder httpClientBuilder = HttpClient.newBuilder() + .connectTimeout(Duration.ofMillis(DEFAULT_CONNECT_TIMEOUT_MILLIS)) + .sslContext(SSLContext.getDefault()); + + if (httpClientConfigCallback != null) { + httpClientBuilder = httpClientConfigCallback.customizeHttpClient(httpClientBuilder); + } + + final HttpClient.Builder finalBuilder = httpClientBuilder; + return AccessController.doPrivileged((PrivilegedAction) finalBuilder::build); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("could not create the default ssl context", e); + } + } + + /** + * Callback used to customize the {@link HttpClient} instance used by a {@link RestHttpClient} instance. + */ + public interface HttpClientConfigCallback { + /** + * Allows to customize the {@link HttpClient} being created and used by the {@link RestHttpClient}. + * Commonly used to customize the default {@link Authenticator} for authentication for communication + * through TLS/SSL without losing any other useful default value that the {@link RestHttpClientBuilder} internally + * sets, like connection pooling. + * + * @param httpClientBuilder the {@link HttpClient.Builder} for customizing the client instance. + */ + HttpClient.Builder customizeHttpClient(HttpClient.Builder httpClientBuilder); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java new file mode 100644 index 0000000000000..120b042c01357 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StatusLine.java @@ -0,0 +1,38 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpClient.Version; +import java.net.http.HttpResponse; +import java.util.Objects; + +/** + * Response status line (protocol, status code) + * Note: This is an experimental API. + */ +public record StatusLine(Version protoVersion, int statusCode) { + /** + * Creates a new status line with the given version and status. + * + * @param protoVersion the protocol version of the response + * @param statusCode the status code of the response + */ + public StatusLine { + protoVersion = protoVersion != null ? protoVersion : Version.HTTP_1_1; + } + + /** + * Creates a new status line from the response + * + * @param response HTTP response + */ + public StatusLine(final HttpResponse response) { + this(Objects.requireNonNull(response, "Response").version(), Objects.requireNonNull(response, "Response").statusCode()); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java new file mode 100644 index 0000000000000..69c943a6123f3 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingRequest.java @@ -0,0 +1,212 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.nio.ByteBuffer; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import org.reactivestreams.Publisher; + +import static java.util.Collections.unmodifiableMap; + +/** + * HTTP Streaming Request to OpenSearch. + * Note: This is an experimental API. + */ +public final class StreamingRequest { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private final RequestOptions options; + private final Publisher publisher; + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param publisher publisher + */ + private StreamingRequest(String method, String endpoint, Publisher publisher) { + this(method, endpoint, Map.of(), publisher, RequestOptions.DEFAULT); + } + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param builder request options builder + * @param publisher publisher + */ + private StreamingRequest( + String method, + String endpoint, + Map parameters, + Publisher publisher, + RequestOptions.Builder builder + ) { + this(method, endpoint, parameters, publisher, builder.build()); + } + + /** + * Constructor + * @param method method + * @param endpoint endpoint + * @param options request options + * @param publisher publisher + */ + private StreamingRequest( + String method, + String endpoint, + Map parameters, + Publisher publisher, + RequestOptions options + ) { + this.method = method; + this.endpoint = endpoint; + this.publisher = publisher; + this.options = options; + } + + public static StreamingRequest.Builder newRequest(String method, String endpoint, Publisher publisher) { + return new StreamingRequest.Builder(method, endpoint, publisher); + } + + /** + * Get endpoint + * @return endpoint + */ + public String endpoint() { + return endpoint; + } + + /** + * Get method + * @return method + */ + public String method() { + return method; + } + + /** + * Get options + * @return options + */ + public RequestOptions options() { + return options; + } + + /** + * Get parameters + * @return parameters + */ + public Map parameters() { + if (options.getParameters().isEmpty()) { + return unmodifiableMap(parameters); + } else { + Map combinedParameters = new HashMap<>(parameters); + combinedParameters.putAll(options.getParameters()); + return unmodifiableMap(combinedParameters); + } + } + + /** + * Body publisher + * @return body publisher + */ + public Publisher body() { + return publisher; + } + + public static final class Builder { + private final String method; + private final String endpoint; + private final Map parameters = new HashMap<>(); + private final Publisher publisher; + private RequestOptions options = RequestOptions.DEFAULT; + + private Builder(String method, String endpoint, Publisher publisher) { + this.method = Objects.requireNonNull(method, "method cannot be null"); + this.endpoint = Objects.requireNonNull(endpoint, "endpoint cannot be null"); + this.publisher = Objects.requireNonNull(publisher, "publisher cannot be null"); + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public void setOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + } + + /** + * Add a query string parameter. + * @param name the name of the url parameter. Must not be null. + * @param value the value of the url url parameter. If {@code null} then + * the parameter is sent as {@code name} rather than {@code name=value} + * @throws IllegalArgumentException if a parameter with that name has + * already been set + */ + public Builder withParameter(String name, String value) { + Objects.requireNonNull(name, "url parameter name cannot be null"); + if (parameters.containsKey(name)) { + throw new IllegalArgumentException("url parameter [" + name + "] has already been set to [" + parameters.get(name) + "]"); + } else { + parameters.put(name, value); + } + return this; + } + + /** + * Add query parameters using the provided map of key value pairs. + * + * @param paramSource a map of key value pairs where the key is the url parameter. + * @throws IllegalArgumentException if a parameter with that name has already been set. + */ + public Builder withParameters(Map paramSource) { + paramSource.forEach(this::withParameter); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions.Builder options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options.build(); + return this; + } + + /** + * Set the portion of an HTTP request to OpenSearch that can be + * manipulated without changing OpenSearch's behavior. + * + * @param options the options to be set. + * @throws NullPointerException if {@code options} is null. + */ + public Builder withOptions(RequestOptions options) { + Objects.requireNonNull(options, "options cannot be null"); + this.options = options; + return this; + } + + public StreamingRequest build() { + return new StreamingRequest(method, endpoint, parameters, publisher, options); + } + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java new file mode 100644 index 0000000000000..629adf2a4b532 --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/StreamingResponse.java @@ -0,0 +1,121 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import java.net.http.HttpHeaders; +import java.net.http.HttpResponse; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.Flow; + +import org.reactivestreams.Publisher; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Mono; + +/** + * HTTP Streaming Response from OpenSearch. + * Note: This is an experimental API. + */ +public final class StreamingResponse { + private final RequestLine requestLine; + private final Mono>>> publisher; + + /** + * Constructor + * @param requestLine request line + * @param publisher message publisher(response with a body) + */ + public StreamingResponse(RequestLine requestLine, Publisher>>> publisher) { + this.requestLine = requestLine; + // We cache the publisher here so the body or / and HttpResponse could + // be consumed independently or/and more than once. + this.publisher = Mono.from(publisher).cache(); + } + + /** + * Get request line + * @return request line + */ + public RequestLine requestLine() { + return requestLine; + } + + /** + * Get response boby {@link Publisher} + * @return response boby {@link Publisher} + */ + public Publisher body() { + return publisher.flatMapMany(m -> { + final boolean compressed = m.headers() + .firstValue("Content-Encoding") + .filter("gzip"::equalsIgnoreCase) + .map(h -> true) + .orElse(false); + return JdkFlowAdapter.flowPublisherToFlux(m.body()).flatMapIterable(t -> t).map(b -> { + if (compressed) { + return BodyUtils.decompress(b); + } else { + return b; + } + }); + }); + } + + /** + * Returns the status line of the current response + */ + @SuppressWarnings("unchecked") + public StatusLine statusLine() { + return new StatusLine( + publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).block() + ); + } + + /** + * Returns a list of all warning headers returned in the response. + */ + @SuppressWarnings("unchecked") + public List warnings() { + return ResponseWarningsExtractor.getWarnings( + publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).block() + ); + } + + /** + * Returns a list of all headers returned in the response. + */ + @SuppressWarnings("unchecked") + public HttpHeaders headers() { + return publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).map(HttpResponse::headers).block(); + } + + /** + * Returns the value of the first header with a specified name of this message. + * If there is more than one matching header in the message the first element is returned. + * If there is no matching header in the message null is returned. + * + * @param name header name + */ + @SuppressWarnings("unchecked") + public String header(String name) { + return publisher.onErrorResume( + ResponseException.class, + e -> Mono.just((HttpResponse>>) e.getResponse().httpResponse()) + ).mapNotNull(response -> response.headers().firstValue(name).orElse(null)).block(); + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java new file mode 100644 index 0000000000000..53a0a1f97fe4b --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningFailureException.java @@ -0,0 +1,79 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.io.IOException; + +import static org.opensearch.internal.httpclient.ResponseException.buildMessage; + +/** + * This exception is used to indicate that one or more {@link Response#warnings()} exist + * and is typically used when the {@link RestHttpClient} is set to fail by setting + * {@link RestHttpClientBuilder#setStrictDeprecationMode(boolean)} to `true`. + */ +// This class extends RuntimeException in order to deal with wrapping that is done in FutureUtils on exception. +// if the exception is not of type OpenSearchException or RuntimeException it will be wrapped in a UncategorizedExecutionException +public final class WarningFailureException extends RuntimeException { + + private final Response response; + + /** + * Creates a {@link WarningFailureException} instance. + * + * @param response the response that contains warnings. + * @throws IOException if there is a problem building the exception message. + */ + public WarningFailureException(Response response) throws IOException { + super(buildMessage(response)); + this.response = response; + } + + /** + * Wrap a {@linkplain WarningFailureException} with another one with the current + * stack trace. This is used during synchronous calls so that the caller + * ends up in the stack trace of the exception thrown. + * + * @param e the exception to be wrapped. + */ + WarningFailureException(WarningFailureException e) { + super(e.getMessage(), e); + this.response = e.getResponse(); + } + + /** + * Returns the {@link Response} that caused this exception to be thrown. + */ + public Response getResponse() { + return response; + } +} diff --git a/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java new file mode 100644 index 0000000000000..1d59a185a6a3e --- /dev/null +++ b/client/rest-http-client/src/main/java/org/opensearch/internal/httpclient/WarningsHandler.java @@ -0,0 +1,80 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.List; + +/** + * Called if there are warnings to determine if those warnings should fail the + * request. + */ +public interface WarningsHandler { + + /** + * Determines whether the given list of warnings should fail the request. + * + * @param warnings a list of warnings. + * @return boolean indicating if the request should fail. + */ + boolean warningsShouldFailRequest(List warnings); + + /** + * The permissive warnings handler. Warnings will not fail the request. + */ + WarningsHandler PERMISSIVE = new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false; + } + + @Override + public String toString() { + return "permissive"; + } + }; + + /** + * The strict warnings handler. Warnings will fail the request. + */ + WarningsHandler STRICT = new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false == warnings.isEmpty(); + } + + @Override + public String toString() { + return "strict"; + } + }; +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java new file mode 100644 index 0000000000000..814cf4bfaff11 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/BouncyCastleThreadFilter.java @@ -0,0 +1,25 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.ThreadFilter; + +/** + * ThreadFilter to exclude ThreadLeak checks for BC’s global background threads + * + *

clone from the original, which is located in ':test:framework'

+ */ +public class BouncyCastleThreadFilter implements ThreadFilter { + @Override + public boolean reject(Thread t) { + String n = t.getName(); + // Ignore BC’s global background threads + return "BC Disposal Daemon".equals(n) || "BC Cleanup Executor".equals(n); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java new file mode 100644 index 0000000000000..c18c6539f8bc4 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HostsTrackingFailureListener.java @@ -0,0 +1,71 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; + +/** + * {@link RestHttpClient.FailureListener} impl that allows to track when it gets called for which host. + */ +class HostsTrackingFailureListener extends RestHttpClient.FailureListener { + private volatile Set hosts = new HashSet<>(); + + @Override + public void onFailure(Node node) { + hosts.add(node.getHost()); + } + + void assertCalled(List nodes) { + HttpHost[] hosts = new HttpHost[nodes.size()]; + for (int i = 0; i < nodes.size(); i++) { + hosts[i] = nodes.get(i).getHost(); + } + assertCalled(hosts); + } + + void assertCalled(HttpHost... hosts) { + assertEquals(hosts.length, this.hosts.size()); + assertThat(this.hosts, containsInAnyOrder(hosts)); + this.hosts.clear(); + } + + void assertNotCalled() { + assertEquals(0, hosts.size()); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java new file mode 100644 index 0000000000000..9ca72d417775c --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/HttpClientThreadLeakFilter.java @@ -0,0 +1,24 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.ThreadFilter; + +import java.net.http.HttpClient; + +/** + * The {@link HttpClient} creates own ASYNC pool based of {@code ForkJoin.commonPool()} which + * is impossible to supress or override. + */ +public final class HttpClientThreadLeakFilter implements ThreadFilter { + @Override + public boolean reject(Thread t) { + return t.getName().startsWith("ForkJoinPool.commonPool-"); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java new file mode 100644 index 0000000000000..ac1dbcb4b8a2e --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestClientTestUtil.java @@ -0,0 +1,120 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.generators.RandomNumbers; +import com.carrotsearch.randomizedtesting.generators.RandomPicks; +import com.carrotsearch.randomizedtesting.generators.RandomStrings; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; + +final class RestClientTestUtil { + + private static final String[] HTTP_METHODS = new String[] { "DELETE", "HEAD", "GET", "OPTIONS", "PATCH", "POST", "PUT", "TRACE" }; + private static final List ALL_STATUS_CODES; + private static final List OK_STATUS_CODES = Arrays.asList(200, 201); + private static final List ALL_ERROR_STATUS_CODES; + private static List ERROR_NO_RETRY_STATUS_CODES = Arrays.asList(400, 401, 403, 404, 405, 500); + private static List ERROR_RETRY_STATUS_CODES = Arrays.asList(502, 503, 504); + + static { + ALL_ERROR_STATUS_CODES = new ArrayList<>(ERROR_RETRY_STATUS_CODES); + ALL_ERROR_STATUS_CODES.addAll(ERROR_NO_RETRY_STATUS_CODES); + ALL_STATUS_CODES = new ArrayList<>(ALL_ERROR_STATUS_CODES); + ALL_STATUS_CODES.addAll(OK_STATUS_CODES); + } + + private RestClientTestUtil() { + + } + + static String[] getHttpMethods() { + return HTTP_METHODS; + } + + static String randomHttpMethod(Random random) { + return RandomPicks.randomFrom(random, HTTP_METHODS); + } + + static int randomStatusCode(Random random) { + return RandomPicks.randomFrom(random, ALL_STATUS_CODES); + } + + static int randomOkStatusCode(Random random) { + return RandomPicks.randomFrom(random, OK_STATUS_CODES); + } + + static int randomErrorNoRetryStatusCode(Random random) { + return RandomPicks.randomFrom(random, ERROR_NO_RETRY_STATUS_CODES); + } + + static int randomErrorRetryStatusCode(Random random) { + return RandomPicks.randomFrom(random, ERROR_RETRY_STATUS_CODES); + } + + static List getOkStatusCodes() { + return OK_STATUS_CODES; + } + + static List getAllErrorStatusCodes() { + return ALL_ERROR_STATUS_CODES; + } + + static List getAllStatusCodes() { + return ALL_STATUS_CODES; + } + + /** + * Create a random number of HTTP headers. + * Generated header names will either be the {@code baseName} plus its index, or exactly the provided {@code baseName} so that the + * we test also support for multiple headers with same key and different values. + */ + static Map> randomHeaders(Random random, final String baseName) { + final Map> headers = new HashMap<>(); + int numHeaders = RandomNumbers.randomIntBetween(random, 0, 5); + for (int i = 0; i < numHeaders; i++) { + String headerName = baseName; + // randomly exercise the code path that supports multiple headers with same key + if (random.nextBoolean()) { + headerName = headerName + i; + } + headers.put(headerName, List.of(RandomStrings.randomAsciiLettersOfLengthBetween(random, 3, 10))); + } + return headers; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java new file mode 100644 index 0000000000000..6a74a752fddd0 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientCompressionTests.java @@ -0,0 +1,142 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient.Version; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +public class RestHttpClientCompressionTests extends RestHttpClientTestCase { + + private static HttpServer httpServer; + + @BeforeClass + public static void startHttpServer() throws Exception { + httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.createContext("/", new GzipResponseHandler()); + httpServer.start(); + } + + @AfterClass + public static void stopHttpServers() throws IOException { + httpServer.stop(0); + httpServer = null; + } + + /** + * A response handler that accepts gzip-encoded data and replies request and response encoding values + * followed by the request body. The response is compressed if "Accept-Encoding" is "gzip". + */ + private static class GzipResponseHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + + // Decode body (if any) + String contentEncoding = exchange.getRequestHeaders().getFirst("Content-Encoding"); + String contentLength = exchange.getRequestHeaders().getFirst("Content-Length"); + InputStream body = exchange.getRequestBody(); + boolean compressedRequest = false; + if ("gzip".equals(contentEncoding)) { + body = new GZIPInputStream(body); + compressedRequest = true; + } + byte[] bytes = body.readAllBytes(); + boolean compress = "gzip".equals(exchange.getRequestHeaders().getFirst("Accept-Encoding")); + if (compress) { + exchange.getResponseHeaders().add("Content-Encoding", "gzip"); + } + + exchange.sendResponseHeaders(200, 0); + + // Encode response if needed + OutputStream out = exchange.getResponseBody(); + if (compress) { + out = new GZIPOutputStream(out); + } + + // Outputs ## + out.write(String.valueOf(contentEncoding).getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write((compress ? "gzip" : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(((compressedRequest == true && contentLength != null) ? contentLength : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(bytes); + out.close(); + + exchange.close(); + } + } + + private RestHttpClient createClient(boolean enableCompression) { + InetSocketAddress address = httpServer.getAddress(); + return RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(enableCompression) + .setHttpClientConfigCallback(builder -> builder.version(Version.HTTP_1_1)) + .build(); + } + + public void testCompressingClientWithContentLengthSync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + // Content-Encoding#Accept-Encoding#Content-Length#Content + // With HttpClient, we don't sent Content-Length so it is always null + Assert.assertEquals("gzip#gzip#null#compressing client", content); + } + } + + public void testCompressingClientContentLengthAsync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + // Server should report it had a compressed request and sent back a compressed response + String content = BodyUtils.getBodyAsString(response.entity()); + + // Content-Encoding#Accept-Encoding#Content-Length#Content + // With HttpClient, we don't sent Content-Length so it is always null + Assert.assertEquals("gzip#gzip#null#compressing client", content); + } + } + + public static class FutureResponse extends CompletableFuture implements ResponseListener { + @Override + public void onSuccess(Response response) { + this.complete(response); + } + + @Override + public void onFailure(Exception exception) { + this.completeExceptionally(exception); + } + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java new file mode 100644 index 0000000000000..64a7bb01f2b7a --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientGzipCompressionTests.java @@ -0,0 +1,191 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.zip.GZIPInputStream; +import java.util.zip.GZIPOutputStream; + +public class RestHttpClientGzipCompressionTests extends RestHttpClientTestCase { + + private static HttpServer httpServer; + + @BeforeClass + public static void startHttpServer() throws Exception { + httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.createContext("/", new GzipResponseHandler()); + httpServer.start(); + } + + @AfterClass + public static void stopHttpServers() throws IOException { + httpServer.stop(0); + httpServer = null; + } + + /** + * A response handler that accepts gzip-encoded data and replies request and response encoding values + * followed by the request body. The response is compressed if "Accept-Encoding" is "gzip". + */ + private static class GzipResponseHandler implements HttpHandler { + @Override + public void handle(HttpExchange exchange) throws IOException { + + // Decode body (if any) + String contentEncoding = exchange.getRequestHeaders().getFirst("Content-Encoding"); + InputStream body = exchange.getRequestBody(); + if ("gzip".equals(contentEncoding)) { + body = new GZIPInputStream(body); + } + byte[] bytes = body.readAllBytes(); + boolean compress = "gzip".equals(exchange.getRequestHeaders().getFirst("Accept-Encoding")); + if (compress) { + exchange.getResponseHeaders().add("Content-Encoding", "gzip"); + } + + exchange.sendResponseHeaders(200, 0); + + // Encode response if needed + OutputStream out = exchange.getResponseBody(); + if (compress) { + out = new GZIPOutputStream(out); + } + + // Outputs ## + out.write(String.valueOf(contentEncoding).getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write((compress ? "gzip" : "null").getBytes(StandardCharsets.UTF_8)); + out.write('#'); + out.write(bytes); + out.close(); + + exchange.close(); + } + } + + private RestHttpClient createClient(boolean enableCompression) { + InetSocketAddress address = httpServer.getAddress(); + return RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(enableCompression) + .build(); + } + + public void testGzipHeaderSync() throws Exception { + try (RestHttpClient restClient = createClient(false)) { + // Send non-compressed request, expect compressed response + Request request = Request.newRequest("POST", "/") + .withEntity(BodyPublishers.ofString("plain request, gzip response")) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeader("Accept-Encoding", "gzip")) + .build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("null#gzip#plain request, gzip response", content); + } + } + + public void testGzipHeaderAsync() throws Exception { + try (RestHttpClient restClient = createClient(false)) { + // Send non-compressed request, expect compressed response + Request request = Request.newRequest("POST", "/") + .withEntity(BodyPublishers.ofString("plain request, gzip response")) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeader("Accept-Encoding", "gzip")) + .build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("null#gzip#plain request, gzip response", content); + } + } + + public void testCompressingClientSync() throws Exception { + try (RestHttpClient restClient = createClient(true)) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + Response response = restClient.performRequest(request); + + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("gzip#gzip#compressing client", content); + } + } + + public void testCompressingClientAsync() throws Exception { + InetSocketAddress address = httpServer.getAddress(); + try ( + RestHttpClient restClient = RestHttpClient.builder(new HttpHost("http", address.getHostString(), address.getPort())) + .setCompressionEnabled(true) + .build() + ) { + Request request = Request.newRequest("POST", "/").withEntity(BodyPublishers.ofString("compressing client")).build(); + + FutureResponse futureResponse = new FutureResponse(); + restClient.performRequestAsync(request, futureResponse); + Response response = futureResponse.get(); + + // Server should report it had a compressed request and sent back a compressed response + String content = BodyUtils.getBodyAsString(response.entity()); + Assert.assertEquals("gzip#gzip#compressing client", content); + } + } + + public static class FutureResponse extends CompletableFuture implements ResponseListener { + @Override + public void onSuccess(Response response) { + this.complete(response); + } + + @Override + public void onFailure(Exception exception) { + this.completeExceptionally(exception); + } + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java new file mode 100644 index 0000000000000..74799a6c5159b --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsIntegTests.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Ignore; + +import java.io.IOException; +import java.net.ConnectException; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import static org.opensearch.internal.httpclient.RestClientTestUtil.getAllStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorNoRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomOkStatusCode; +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against real http servers, multiple hosts. Also tests failover by randomly shutting down hosts. + */ +public class RestHttpClientMultipleHostsIntegTests extends RestHttpClientTestCase { + + private static WaitForCancelHandler waitForCancelHandler; + private static HttpServer[] httpServers; + private static HttpHost[] httpHosts; + private static boolean stoppedFirstHost = false; + private static String pathPrefixWithoutLeadingSlash; + private static String pathPrefix; + private static RestHttpClient restClient; + + @BeforeClass + public static void startHttpServer() throws Exception { + if (randomBoolean()) { + pathPrefixWithoutLeadingSlash = "testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5); + pathPrefix = "/" + pathPrefixWithoutLeadingSlash; + } else { + pathPrefix = pathPrefixWithoutLeadingSlash = ""; + } + int numHttpServers = randomIntBetween(2, 4); + httpServers = new HttpServer[numHttpServers]; + httpHosts = new HttpHost[numHttpServers]; + waitForCancelHandler = new WaitForCancelHandler(); + for (int i = 0; i < numHttpServers; i++) { + HttpServer httpServer = createHttpServer(); + httpServers[i] = httpServer; + httpHosts[i] = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + } + restClient = buildRestClient(NodeSelector.ANY); + } + + private static RestHttpClient buildRestClient(NodeSelector nodeSelector) { + RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHosts); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix((randomBoolean() ? "/" : "") + pathPrefixWithoutLeadingSlash); + } + restClientBuilder.setNodeSelector(nodeSelector); + return restClientBuilder.build(); + } + + private static HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + httpServer.createContext(pathPrefix + "/wait", waitForCancelHandler); + return httpServer; + } + + private static class WaitForCancelHandler implements HttpHandler { + private volatile CountDownLatch requestCameInLatch; + private volatile CountDownLatch cancelHandlerLatch; + + void reset() { + cancelHandlerLatch = new CountDownLatch(1); + requestCameInLatch = new CountDownLatch(1); + } + + void cancelDone() { + cancelHandlerLatch.countDown(); + } + + void awaitRequest() throws InterruptedException { + requestCameInLatch.await(); + } + + @Override + public void handle(HttpExchange exchange) throws IOException { + requestCameInLatch.countDown(); + try { + cancelHandlerLatch.await(); + } catch (InterruptedException ignore) {} finally { + exchange.sendResponseHeaders(200, 0); + exchange.close(); + } + } + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + httpExchange.getRequestBody().close(); + httpExchange.sendResponseHeaders(statusCode, -1); + httpExchange.close(); + } + } + + @AfterClass + public static void stopHttpServers() throws IOException { + restClient.close(); + restClient = null; + for (HttpServer httpServer : httpServers) { + httpServer.stop(0); + } + httpServers = null; + } + + @Before + public void stopRandomHost() { + // verify that shutting down some hosts doesn't matter as long as one working host is left behind + if (httpServers.length > 1 && randomBoolean()) { + List updatedHttpServers = new ArrayList<>(httpServers.length - 1); + int nodeIndex = randomIntBetween(0, httpServers.length - 1); + if (0 == nodeIndex) { + stoppedFirstHost = true; + } + for (int i = 0; i < httpServers.length; i++) { + HttpServer httpServer = httpServers[i]; + if (i == nodeIndex) { + httpServer.stop(0); + } else { + updatedHttpServers.add(httpServer); + } + } + httpServers = updatedHttpServers.toArray(new HttpServer[0]); + } + } + + public void testSyncRequests() throws IOException { + int numRequests = randomIntBetween(5, 20); + for (int i = 0; i < numRequests; i++) { + final String method = RestClientTestUtil.randomHttpMethod(getRandom()); + // we don't test status codes that are subject to retries as they interfere with hosts being stopped + final int statusCode = randomBoolean() ? randomOkStatusCode(getRandom()) : randomErrorNoRetryStatusCode(getRandom()); + Response response; + try { + response = restClient.performRequest(Request.newRequest(method, "/" + statusCode).build()); + } catch (ResponseException responseException) { + response = responseException.getResponse(); + } + assertEquals(method, response.requestLine().method()); + assertEquals(statusCode, response.statusLine().statusCode()); + assertEquals((pathPrefix.length() > 0 ? pathPrefix : "") + "/" + statusCode, response.requestLine().uri()); + } + } + + public void testAsyncRequests() throws Exception { + int numRequests = randomIntBetween(5, 20); + final CountDownLatch latch = new CountDownLatch(numRequests); + final List responses = new CopyOnWriteArrayList<>(); + for (int i = 0; i < numRequests; i++) { + final String method = RestClientTestUtil.randomHttpMethod(getRandom()); + // we don't test status codes that are subject to retries as they interfere with hosts being stopped + final int statusCode = randomBoolean() ? randomOkStatusCode(getRandom()) : randomErrorNoRetryStatusCode(getRandom()); + restClient.performRequestAsync(Request.newRequest(method, "/" + statusCode).build(), new ResponseListener() { + @Override + public void onSuccess(Response response) { + responses.add(new TestResponse(method, statusCode, response)); + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + responses.add(new TestResponse(method, statusCode, exception)); + latch.countDown(); + } + }); + } + assertTrue(latch.await(5, TimeUnit.SECONDS)); + + assertEquals(numRequests, responses.size()); + for (TestResponse testResponse : responses) { + Response response = testResponse.getResponse(); + assertEquals(testResponse.method, response.requestLine().method()); + assertEquals(testResponse.statusCode, response.statusLine().statusCode()); + assertEquals((pathPrefix.length() > 0 ? pathPrefix : "") + "/" + testResponse.statusCode, response.requestLine().uri()); + } + } + + @Ignore("https://github.com/elastic/elasticsearch/issues/45577") + public void testCancelAsyncRequests() throws Exception { + int numRequests = randomIntBetween(5, 20); + final List responses = new CopyOnWriteArrayList<>(); + final List exceptions = new CopyOnWriteArrayList<>(); + for (int i = 0; i < numRequests; i++) { + CountDownLatch latch = new CountDownLatch(1); + waitForCancelHandler.reset(); + Cancellable cancellable = restClient.performRequestAsync(Request.newRequest("GET", "/wait").build(), new ResponseListener() { + @Override + public void onSuccess(Response response) { + responses.add(response); + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + exceptions.add(exception); + latch.countDown(); + } + }); + if (randomBoolean()) { + // we wait for the request to get to the server-side otherwise we almost always cancel + // the request artificially on the client-side before even sending it + waitForCancelHandler.awaitRequest(); + } + cancellable.cancel(); + waitForCancelHandler.cancelDone(); + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } + assertEquals(0, responses.size()); + assertEquals(numRequests, exceptions.size()); + for (Exception exception : exceptions) { + assertThat(exception, instanceOf(CancellationException.class)); + } + } + + /** + * Test host selector against a real server and + * test what happens after calling + */ + public void testNodeSelector() throws Exception { + try (RestHttpClient restClient = buildRestClient(firstPositionNodeSelector())) { + Request request = Request.newRequest("GET", "/200").build(); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + if (stoppedFirstHost) { + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + fail("expected to fail to connect"); + } catch (ConnectException e) { + // HttpClient isn't consistent here. Sometimes the message is even null! + if (e.getMessage() != null) { + assertThat(e.getMessage(), containsString("Connection refused")); + } + } + } else { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(httpHosts[0], response.host()); + } + } + } + } + + private static class TestResponse { + private final String method; + private final int statusCode; + private final Object response; + + TestResponse(String method, int statusCode, Object response) { + this.method = method; + this.statusCode = statusCode; + this.response = response; + } + + Response getResponse() { + if (response instanceof Response) { + return (Response) response; + } + if (response instanceof ResponseException) { + return ((ResponseException) response).getResponse(); + } + throw new AssertionError("unexpected response " + response.getClass()); + } + } + + private NodeSelector firstPositionNodeSelector() { + return nodes -> { + for (Iterator itr = nodes.iterator(); itr.hasNext();) { + if (httpHosts[0] != itr.next().getHost()) { + itr.remove(); + } + } + }; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java new file mode 100644 index 0000000000000..2e644f427c4ee --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientMultipleHostsTests.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.generators.RandomNumbers; + +import org.junit.After; + +import java.io.IOException; +import java.net.http.HttpClient; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; + +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorNoRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomErrorRetryStatusCode; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomHttpMethod; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomOkStatusCode; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * Tests for {@link RestHttpClient} behaviour against multiple hosts: fail-over, denylisting etc. + * Relies on a mock http client to intercept requests and return desired responses based on request path. + */ +public class RestHttpClientMultipleHostsTests extends RestHttpClientTestCase { + + private ExecutorService exec = Executors.newFixedThreadPool(1); + private List nodes; + private HostsTrackingFailureListener failureListener; + + public RestHttpClient createRestClient(NodeSelector nodeSelector) { + HttpClient httpClient = RestHttpClientSingleHostTests.mockHttpClient(exec); + int numNodes = RandomNumbers.randomIntBetween(getRandom(), 2, 5); + nodes = new ArrayList<>(numNodes); + for (int i = 0; i < numNodes; i++) { + nodes.add(new Node(new HttpHost("http", "localhost", 9200 + i))); + } + nodes = Collections.unmodifiableList(nodes); + failureListener = new HostsTrackingFailureListener(); + return new RestHttpClient(httpClient, Map.of(), nodes, null, failureListener, nodeSelector, false, false); + } + + /** + * Shutdown the executor so we don't leak threads into other test runs. + */ + @After + public void shutdownExec() { + exec.shutdown(); + } + + public void testRoundRobinOkStatusCodes() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + int numIters = RandomNumbers.randomIntBetween(getRandom(), 1, 5); + for (int i = 0; i < numIters; i++) { + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + int statusCode = randomOkStatusCode(getRandom()); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), "/" + statusCode).build() + ); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + failureListener.assertNotCalled(); + } + + public void testRoundRobinNoRetryErrors() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + int numIters = RandomNumbers.randomIntBetween(getRandom(), 1, 5); + for (int i = 0; i < numIters; i++) { + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + String method = randomHttpMethod(getRandom()); + int statusCode = randomErrorNoRetryStatusCode(getRandom()); + try { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(method, "/" + statusCode).build() + ); + if (method.equals("HEAD") && statusCode == 404) { + // no exception gets thrown although we got a 404 + assertEquals(404, response.statusLine().statusCode()); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + } else { + fail("request should have failed"); + } + } catch (ResponseException e) { + if (method.equals("HEAD") && statusCode == 404) { + throw e; + } + Response response = e.getResponse(); + assertEquals(statusCode, response.statusLine().statusCode()); + assertTrue("host not found: " + response.host(), hostsSet.remove(response.host())); + assertEquals(0, e.getSuppressed().length); + } + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + failureListener.assertNotCalled(); + } + + public void testRoundRobinRetryErrors() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.ANY); + String retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Set hostsSet = hostsSet(); + // first request causes all the hosts to be denylisted, the returned exception holds one suppressed exception each + failureListener.assertCalled(nodes); + do { + Response response = e.getResponse(); + assertEquals(Integer.parseInt(retryEndpoint.substring(1)), response.statusLine().statusCode()); + assertTrue("host [" + response.host() + "] not found, most likely used multiple times", hostsSet.remove(response.host())); + if (e.getSuppressed().length > 0) { + assertEquals(1, e.getSuppressed().length); + Throwable suppressed = e.getSuppressed()[0]; + assertThat(suppressed, instanceOf(ResponseException.class)); + e = (ResponseException) suppressed; + } else { + e = null; + } + } while (e != null); + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } catch (IOException e) { + Set hostsSet = hostsSet(); + // first request causes all the hosts to be denylisted, the returned exception holds one suppressed exception each + failureListener.assertCalled(nodes); + do { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertTrue("host [" + httpHost + "] not found, most likely used multiple times", hostsSet.remove(httpHost)); + if (e.getSuppressed().length > 0) { + assertEquals(1, e.getSuppressed().length); + Throwable suppressed = e.getSuppressed()[0]; + assertThat(suppressed, instanceOf(IOException.class)); + e = (IOException) suppressed; + } else { + e = null; + } + } while (e != null); + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + } + + int numIters = RandomNumbers.randomIntBetween(getRandom(), 2, 5); + for (int i = 1; i <= numIters; i++) { + // check that one different host is resurrected at each new attempt + Set hostsSet = hostsSet(); + for (int j = 0; j < nodes.size(); j++) { + retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(Integer.parseInt(retryEndpoint.substring(1)))); + assertTrue( + "host [" + response.host() + "] not found, most likely used multiple times", + hostsSet.remove(response.host()) + ); + // after the first request, all hosts are denylisted, a single one gets resurrected each time + failureListener.assertCalled(response.host()); + assertEquals(0, e.getSuppressed().length); + } catch (IOException e) { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertTrue("host [" + httpHost + "] not found, most likely used multiple times", hostsSet.remove(httpHost)); + // after the first request, all hosts are denylisted, a single one gets resurrected each time + failureListener.assertCalled(httpHost); + assertEquals(0, e.getSuppressed().length); + } + } + assertEquals("every host should have been used but some weren't: " + hostsSet, 0, hostsSet.size()); + if (getRandom().nextBoolean()) { + // mark one host back alive through a successful request and check that all requests after that are sent to it + HttpHost selectedHost = null; + int iters = RandomNumbers.randomIntBetween(getRandom(), 2, 10); + for (int y = 0; y < iters; y++) { + int statusCode = randomErrorNoRetryStatusCode(getRandom()); + Response response; + try { + response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), "/" + statusCode).build() + ); + } catch (ResponseException e) { + response = e.getResponse(); + } + assertThat(response.statusLine().statusCode(), equalTo(statusCode)); + if (selectedHost == null) { + selectedHost = response.host(); + } else { + assertThat(response.host(), equalTo(selectedHost)); + } + } + failureListener.assertNotCalled(); + // let the selected host catch up on number of failures, it gets selected a consecutive number of times as it's the one + // selected to be retried earlier (due to lower number of failures) till all the hosts have the same number of failures + for (int y = 0; y < i + 1; y++) { + retryEndpoint = randomErrorRetryEndpoint(); + try { + RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest(randomHttpMethod(getRandom()), retryEndpoint).build() + ); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(Integer.parseInt(retryEndpoint.substring(1)))); + assertThat(response.host(), equalTo(selectedHost)); + failureListener.assertCalled(selectedHost); + } catch (IOException e) { + HttpHost httpHost = HttpHost.create(e.getMessage()); + assertThat(httpHost, equalTo(selectedHost)); + failureListener.assertCalled(selectedHost); + } + } + } + } + } + + public void testNodeSelector() throws Exception { + NodeSelector firstPositionOnly = restClientNodes -> { + boolean found = false; + for (Iterator itr = restClientNodes.iterator(); itr.hasNext();) { + if (nodes.get(0) == itr.next()) { + found = true; + } else { + itr.remove(); + } + } + assertTrue(found); + }; + RestHttpClient restClient = createRestClient(firstPositionOnly); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + Request request = Request.newRequest("GET", "/200").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(nodes.get(0).getHost(), response.host()); + } + } + + public void testSetNodes() throws Exception { + RestHttpClient restClient = createRestClient(NodeSelector.SKIP_DEDICATED_CLUSTER_MANAGERS); + List newNodes = new ArrayList<>(nodes.size()); + for (int i = 0; i < nodes.size(); i++) { + Node.Roles roles = i == 0 + ? new Node.Roles(new TreeSet<>(Arrays.asList("data", "ingest"))) + : new Node.Roles(new TreeSet<>(Arrays.asList("master"))); + newNodes.add(new Node(nodes.get(i).getHost(), null, null, null, roles, null)); + } + restClient.setNodes(newNodes); + int rounds = between(1, 10); + for (int i = 0; i < rounds; i++) { + /* + * Run the request more than once to verify that the + * NodeSelector overrides the round robin behavior. + */ + Request request = Request.newRequest("GET", "/200").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(newNodes.get(0).getHost(), response.host()); + } + } + + private static String randomErrorRetryEndpoint() { + switch (RandomNumbers.randomIntBetween(getRandom(), 0, 3)) { + case 0: + return "/" + randomErrorRetryStatusCode(getRandom()); + case 1: + return "/coe"; + case 2: + return "/soe"; + case 3: + return "/ioe"; + } + throw new UnsupportedOperationException(); + } + + /** + * Build a mutable {@link Set} containing all the {@link Node#getHost() hosts} + * in use by the test. + */ + private Set hostsSet() { + Set hosts = new HashSet<>(); + for (Node node : nodes) { + hosts.add(node.getHost()); + } + return hosts; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java new file mode 100644 index 0000000000000..b76081bd684e4 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostIntegTests.java @@ -0,0 +1,501 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.net.Authenticator; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.PasswordAuthentication; +import java.net.http.HttpClient; +import java.net.http.HttpRequest.BodyPublishers; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Base64; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.LockSupport; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.CoreMatchers.startsWith; +import static org.hamcrest.Matchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against a real http server, one single host. + */ +public class RestHttpClientSingleHostIntegTests extends RestHttpClientTestCase { + + private HttpServer httpServer; + private RestHttpClient restClient; + private String pathPrefix; + private Map> defaultHeaders; + private WaitForCancelHandler waitForCancelHandler; + private ExecutorService httpClientExecutor; + private ExecutorService httpServerExecutor; + + @Before + public void startHttpServer() throws Exception { + pathPrefix = randomBoolean() ? "/testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5) : ""; + httpServer = createHttpServer(); + httpClientExecutor = Executors.newFixedThreadPool(5); + httpServerExecutor = Executors.newFixedThreadPool(10); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + restClient = createRestClient(false, true); + } + + private HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.setExecutor(httpServerExecutor); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : RestClientTestUtil.getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + waitForCancelHandler = new WaitForCancelHandler(); + httpServer.createContext(pathPrefix + "/wait", waitForCancelHandler); + return httpServer; + } + + private static class WaitForCancelHandler implements HttpHandler { + + private final CountDownLatch cancelHandlerLatch = new CountDownLatch(1); + + void cancelDone() { + cancelHandlerLatch.countDown(); + } + + @Override + public void handle(HttpExchange exchange) throws IOException { + try { + cancelHandlerLatch.await(); + } catch (InterruptedException ignore) {} finally { + exchange.sendResponseHeaders(200, 0); + exchange.close(); + } + } + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + // copy request body to response body so we can verify it was sent + StringBuilder body = new StringBuilder(); + try (InputStreamReader reader = new InputStreamReader(httpExchange.getRequestBody(), StandardCharsets.UTF_8)) { + char[] buffer = new char[256]; + int read; + while ((read = reader.read(buffer)) != -1) { + body.append(buffer, 0, read); + } + } + // copy request headers to response headers so we can verify they were sent + Headers requestHeaders = httpExchange.getRequestHeaders(); + Headers responseHeaders = httpExchange.getResponseHeaders(); + for (Map.Entry> header : requestHeaders.entrySet()) { + responseHeaders.put(header.getKey(), header.getValue()); + } + httpExchange.getRequestBody().close(); + httpExchange.sendResponseHeaders(statusCode, body.length() == 0 ? -1 : body.length()); + if (body.length() > 0) { + try (OutputStream out = httpExchange.getResponseBody()) { + out.write(body.toString().getBytes(StandardCharsets.UTF_8)); + } + } + httpExchange.close(); + } + } + + private RestHttpClient createRestClient(final boolean useAuth, final boolean usePreemptiveAuth) { + final HttpHost httpHost = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + final RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHost).setDefaultHeaders(defaultHeaders); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix(pathPrefix); + } + + if (useAuth) { + if (usePreemptiveAuth == false) { + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).authenticator(new Authenticator() { + @Override + protected PasswordAuthentication getPasswordAuthentication() { + return new PasswordAuthentication("user", "pass".toCharArray()); + } + }); + } + }); + } else { + final String auth = Base64.getEncoder().encodeToString(("user" + ":" + "pass").getBytes(StandardCharsets.UTF_8)); + final Map> headers = new HashMap<>(defaultHeaders); + headers.put("Authorization", List.of("Basic " + auth)); + restClientBuilder.setDefaultHeaders(headers); + + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor); + } + }); + } + } else { + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).connectTimeout(Duration.ofSeconds(15)); + } + }); + } + + return restClientBuilder.build(); + } + + @After + public void stopHttpServers() throws IOException, InterruptedException { + restClient.close(); + restClient = null; + httpServer.stop(0); + httpServer = null; + + httpServerExecutor.shutdown(); + if (httpServerExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpServerExecutor.shutdownNow(); + } + + httpClientExecutor.shutdown(); + if (httpClientExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpClientExecutor.shutdownNow(); + } + } + + /** + * Tests sending a bunch of async requests works well (e.g. no TimeoutException from the leased pool) + * See https://github.com/elastic/elasticsearch/issues/24069 + */ + public void testManyAsyncRequests() throws Exception { + int iters = randomIntBetween(500, 1000); + final CountDownLatch latch = new CountDownLatch(iters); + final List exceptions = new CopyOnWriteArrayList<>(); + for (int i = 0; i < iters; i++) { + Request request = Request.newRequest("PUT", "/200").withEntity(BodyPublishers.ofString("{}")).build(); + // Add random jitter so HttpServer will not refuse the connections + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(randomLongBetween(1, 5))); + restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + latch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + exceptions.add(exception); + latch.countDown(); + } + }); + } + + assertTrue("timeout waiting for requests to be sent", latch.await(10, TimeUnit.SECONDS)); + if (exceptions.isEmpty() == false) { + AssertionError error = new AssertionError( + "expected no failures but got some. see suppressed for first 10 of [" + exceptions.size() + "] failures" + ); + for (Exception exception : exceptions.subList(0, Math.min(10, exceptions.size()))) { + error.addSuppressed(exception); + } + throw error; + } + } + + public void testCancelAsyncRequest() throws Exception { + Request request = Request.newRequest(RestClientTestUtil.randomHttpMethod(getRandom()), "/wait").build(); + CountDownLatch requestLatch = new CountDownLatch(1); + AtomicReference error = new AtomicReference<>(); + Cancellable cancellable = restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + throw new AssertionError("onResponse called unexpectedly"); + } + + @Override + public void onFailure(Exception exception) { + error.set(exception); + requestLatch.countDown(); + } + }); + cancellable.cancel(); + waitForCancelHandler.cancelDone(); + assertTrue(requestLatch.await(5, TimeUnit.SECONDS)); + assertThat(error.get(), instanceOf(CancellationException.class)); + } + + /** + * End to end test for headers. We test it explicitly against a real http client as there are different ways + * to set/add headers to the {@link HttpClient}. + * Exercises the test http server ability to send back whatever headers it received. + */ + public void testHeaders() throws Exception { + for (String method : RestClientTestUtil.getHttpMethods()) { + final Set standardHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + standardHeaders.addAll( + Arrays.asList("Connection", "Host", "User-agent", "Date", "Upgrade", "HTTP2-Settings", "Content-Length") + ); + + final Map> requestHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + final int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + Request request = Request.newRequest(method, "/" + statusCode) + .withOptions(RequestOptions.DEFAULT.toBuilder().addHeaders(requestHeaders)) + .build(); + Response esResponse; + try { + esResponse = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + + assertHeaders(defaultHeaders, requestHeaders, esResponse.headers(), standardHeaders); + final Set removedHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + for (final Map.Entry> responseHeader : esResponse.headers().map().entrySet()) { + String name = responseHeader.getKey().toLowerCase(Locale.ROOT); + // Some headers could be returned multiple times in response, like Connection fe. + if (name.startsWith("header") == false && removedHeaders.contains(name) == false) { + assertTrue("unknown header was returned " + name, standardHeaders.remove(name)); + removedHeaders.add(name); + } + } + assertTrue("some expected standard headers weren't returned: " + standardHeaders, standardHeaders.isEmpty()); + } + } + + /** + * End to end test for delete with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testDeleteWithBody() throws Exception { + bodyTest("DELETE"); + } + + /** + * End to end test for get with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testGetWithBody() throws Exception { + bodyTest("GET"); + } + + public void testEncodeParams() throws Exception { + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "this/is/the/routing").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=this%2Fis%2Fthe%2Frouting", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "this|is|the|routing").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=this%7Cis%7Cthe%7Crouting", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "routing#1").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=routing%231", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "中文").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=%E4%B8%AD%E6%96%87", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo+bar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo+bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%2Bbar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo/bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%2Fbar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("routing", "foo^bar").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?routing=foo%5Ebar", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("pretty", null).build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty", response.requestLine().uri()); + } + { + Request request = Request.newRequest("PUT", "/200").withParameter("pretty", "").build(); + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty=", response.requestLine().uri()); + } + } + + /** + * Verify that credentials are sent on the first request with preemptive auth enabled (default when provided with credentials). + */ + public void testPreemptiveAuthEnabled() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, true)) { + for (final String method : methods) { + final Response response = bodyTest(restClient, method); + assertThat(response.header("Authorization"), startsWith("Basic")); + } + } + } + + /** + * Verify that credentials are not sent on the first request with preemptive auth disabled. + */ + public void testPreemptiveAuthDisabled() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, false)) { + for (final String method : methods) { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + if (statusCode == 401) { + final IOException ex = assertThrows(IOException.class, () -> bodyTest(restClient, method, statusCode, Map.of())); + assertThat(ex.getMessage(), equalTo("WWW-Authenticate header missing for response code 401")); + } else { + final Response response = bodyTest(restClient, method, statusCode, Map.of()); + assertThat(response.header("Authorization"), nullValue()); + } + } + } + } + + /** + * Verify that credentials continue to be sent even if a 401 (Unauthorized) response is received + */ + public void testAuthCredentialsAreNotClearedOnAuthChallenge() throws Exception { + final String[] methods = { "POST", "PUT", "GET", "DELETE" }; + + try (RestHttpClient restClient = createRestClient(true, true)) { + for (final String method : methods) { + Map> realmHeader = Map.of("WWW-Authenticate", List.of("Basic realm=\"test\"")); + final Response response401 = bodyTest(restClient, method, 401, realmHeader); + assertThat(response401.header("Authorization"), startsWith("Basic")); + + final Response response200 = bodyTest(restClient, method, 200, Map.of()); + assertThat(response200.header("Authorization"), startsWith("Basic")); + } + } + } + + public void testUrlWithoutLeadingSlash() throws Exception { + if (pathPrefix.length() == 0) { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added even if a pathPrefix is not configured (HttpClient uses full URI) + assertEquals(200, response.statusLine().statusCode()); + } else { + { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added if a pathPrefix is configured + assertEquals(200, response.statusLine().statusCode()); + } + { + // pathPrefix is not required to start with '/', will be added automatically + try ( + RestHttpClient restClient = RestHttpClient.builder( + new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()) + ).setPathPrefix(pathPrefix.substring(1)).build() + ) { + Response response = RestHttpClientSingleHostTests.performRequestSyncOrAsync( + restClient, + Request.newRequest("GET", "200").build() + ); + // a trailing slash gets automatically added if a pathPrefix is configured + assertEquals(200, response.statusLine().statusCode()); + } + } + } + } + + private Response bodyTest(final String method) throws Exception { + return bodyTest(restClient, method); + } + + private Response bodyTest(final RestHttpClient restClient, final String method) throws Exception { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + return bodyTest(restClient, method, statusCode, Map.of()); + } + + private Response bodyTest(RestHttpClient restClient, String method, int statusCode, Map> headers) + throws Exception { + String requestBody = "{ \"field\": \"value\" }"; + Request.Builder builder = Request.newRequest(method, "/" + statusCode).withEntity(requestBody); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> options.addHeader(header.getKey(), v)); + } + builder = builder.withOptions(options); + Response esResponse; + try { + esResponse = RestHttpClientSingleHostTests.performRequestSyncOrAsync(restClient, builder.build()); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + assertEquals(requestBody, BodyUtils.getBodyAsString(esResponse)); + + return esResponse; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java new file mode 100644 index 0000000000000..6439826eac37b --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostStreamingIntegTests.java @@ -0,0 +1,180 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.net.http.HttpClient; +import java.nio.CharBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.junit.Assert.assertEquals; + +/** + * Integration test to check interaction between {@link RestHttpClient} and {@link HttpClient}. + * Works against a real http server, one single host, use streaming. + */ +public class RestHttpClientSingleHostStreamingIntegTests extends RestHttpClientTestCase { + + private HttpServer httpServer; + private RestHttpClient restClient; + private String pathPrefix; + private Map> defaultHeaders; + private ExecutorService httpClientExecutor; + + @Before + public void startHttpServer() throws Exception { + pathPrefix = randomBoolean() ? "/testPathPrefix/" + randomAsciiLettersOfLengthBetween(1, 5) : ""; + httpServer = createHttpServer(); + httpClientExecutor = Executors.newWorkStealingPool(); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + restClient = createRestClient(); + } + + private HttpServer createHttpServer() throws Exception { + HttpServer httpServer = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + httpServer.start(); + // returns a different status code depending on the path + for (int statusCode : RestClientTestUtil.getAllStatusCodes()) { + httpServer.createContext(pathPrefix + "/" + statusCode, new ResponseHandler(statusCode)); + } + return httpServer; + } + + private static class ResponseHandler implements HttpHandler { + private final int statusCode; + + ResponseHandler(int statusCode) { + this.statusCode = statusCode; + } + + @Override + public void handle(HttpExchange httpExchange) throws IOException { + Headers requestHeaders = httpExchange.getRequestHeaders(); + Headers responseHeaders = httpExchange.getResponseHeaders(); + for (Map.Entry> header : requestHeaders.entrySet()) { + responseHeaders.put(header.getKey(), header.getValue()); + } + try { + httpExchange.sendResponseHeaders(statusCode, 0); + httpExchange.getRequestBody().transferTo(httpExchange.getResponseBody()); + httpExchange.getResponseBody().flush(); + } finally { + httpExchange.getRequestBody().close(); + httpExchange.getResponseBody().close(); + httpExchange.close(); + } + } + } + + private RestHttpClient createRestClient() { + final HttpHost httpHost = new HttpHost("http", httpServer.getAddress().getHostString(), httpServer.getAddress().getPort()); + final RestHttpClientBuilder restClientBuilder = RestHttpClient.builder(httpHost) + .setDefaultHeaders(defaultHeaders) + .setCompressionEnabled(randomBoolean()); + if (pathPrefix.length() > 0) { + restClientBuilder.setPathPrefix(pathPrefix); + } + + restClientBuilder.setHttpClientConfigCallback(new RestHttpClientBuilder.HttpClientConfigCallback() { + @Override + public HttpClient.Builder customizeHttpClient(final HttpClient.Builder httpClientBuilder) { + return httpClientBuilder.executor(httpClientExecutor).connectTimeout(Duration.ofSeconds(15)); + } + }); + + return restClientBuilder.build(); + } + + @After + public void stopHttpServers() throws IOException, InterruptedException { + restClient.close(); + restClient = null; + httpServer.stop(0); + httpServer = null; + + httpClientExecutor.shutdown(); + if (httpClientExecutor.awaitTermination(30, TimeUnit.SECONDS) == false) { + httpClientExecutor.shutdownNow(); + } + } + + /** + * End to end test for delete with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testDeleteWithBody() throws Exception { + bodyTest("DELETE"); + } + + /** + * End to end test for get with body. We test it explicitly as it is not supported + * out of the box by {@link HttpClient}. + * Exercises the test http server ability to send back whatever body it received. + */ + public void testGetWithBody() throws Exception { + bodyTest("GET"); + } + + private StreamingResponse bodyTest(final String method) throws Exception { + int statusCode = RestClientTestUtil.randomStatusCode(getRandom()); + return bodyTest(restClient, method, statusCode, Map.of()); + } + + private StreamingResponse bodyTest(RestHttpClient restClient, String method, int statusCode, Map> headers) + throws Exception { + String requestBody = "{ \"field\": \"value\" }"; + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> options.addHeader(header.getKey(), v)); + } + + final StreamingRequest request = StreamingRequest.newRequest( + method, + "/" + statusCode, + Mono.just(StandardCharsets.UTF_8.encode(requestBody)) + ).withOptions(options).build(); + StreamingResponse esResponse = restClient.streamRequest(request); + + assertEquals(method, esResponse.requestLine().method()); + assertEquals(statusCode, esResponse.statusLine().statusCode()); + assertEquals(pathPrefix + "/" + statusCode, esResponse.requestLine().uri()); + + if (statusCode >= 200 && statusCode < 400) { + StepVerifier.create(Flux.from(esResponse.body()).map(StandardCharsets.UTF_8::decode).map(CharBuffer::toString)) + .expectNextMatches(s -> s.equals(requestBody)) + .expectComplete() + .verify(Duration.ofSeconds(5)); + } else { + StepVerifier.create(Flux.from(esResponse.body())).expectError(ResponseException.class).verify(Duration.ofSeconds(5)); + } + + return esResponse; + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java new file mode 100644 index 0000000000000..eb8a16d257cc0 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientSingleHostTests.java @@ -0,0 +1,728 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; +import org.junit.After; +import org.junit.Before; + +import javax.net.ssl.SSLContext; +import javax.net.ssl.SSLHandshakeException; +import javax.net.ssl.SSLParameters; +import javax.net.ssl.SSLSession; + +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.net.Authenticator; +import java.net.CookieHandler; +import java.net.ProxySelector; +import java.net.SocketTimeoutException; +import java.net.URI; +import java.net.URISyntaxException; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpClient.Version; +import java.net.http.HttpHeaders; +import java.net.http.HttpRequest; +import java.net.http.HttpRequest.BodyPublisher; +import java.net.http.HttpRequest.BodyPublishers; +import java.net.http.HttpResponse; +import java.net.http.HttpResponse.BodyHandler; +import java.net.http.HttpResponse.PushPromiseHandler; +import java.net.http.HttpTimeoutException; +import java.nio.channels.ClosedChannelException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.TreeSet; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.atomic.LongAdder; +import java.util.stream.Collectors; + +import static java.util.Collections.singletonList; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getAllErrorStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getHttpMethods; +import static org.opensearch.internal.httpclient.RestClientTestUtil.getOkStatusCodes; +import static org.opensearch.internal.httpclient.RestClientTestUtil.randomStatusCode; +import static org.hamcrest.CoreMatchers.containsString; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.fail; + +/** + * Tests for basic functionality of {@link RestHttpClient} against one single host: tests http requests being sent, headers, + * body, different status codes and corresponding responses/exceptions. + * Relies on a mock http client to intercept requests and return desired responses based on request path. + */ +public class RestHttpClientSingleHostTests extends RestHttpClientTestCase { + private static final Log logger = LogFactory.getLog(RestHttpClientSingleHostTests.class); + + private ExecutorService exec = Executors.newFixedThreadPool(1); + private RestHttpClient restClient; + private Map> defaultHeaders; + private Node node; + private HttpClient httpClient; + private HostsTrackingFailureListener failureListener; + private boolean strictDeprecationMode; + private LongAdder requests; + private AtomicReference requestProducerCapture; + + @Before + public void createRestClient() { + requests = new LongAdder(); + requestProducerCapture = new AtomicReference<>(); + httpClient = mockHttpClient(exec, (requestProducer, bodyHandler) -> { + requests.increment(); + requestProducerCapture.set(requestProducer); + }); + defaultHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header-default"); + node = new Node(new HttpHost("http", "localhost", 9200)); + failureListener = new HostsTrackingFailureListener(); + strictDeprecationMode = randomBoolean(); + restClient = new RestHttpClient( + this.httpClient, + defaultHeaders, + singletonList(node), + null, + failureListener, + NodeSelector.ANY, + strictDeprecationMode, + false + ); + } + + interface HttpClientListener { + void onExecute(HttpRequest requestProducer, BodyHandler bodyHandler); + } + + @SuppressWarnings("unchecked") + static HttpClient mockHttpClient(final ExecutorService exec, final HttpClientListener... listeners) { + HttpClient httpClient = new HttpClient() { + @Override + public Optional authenticator() { + return Optional.empty(); + } + + @Override + public Optional connectTimeout() { + return Optional.empty(); + } + + @Override + public Optional cookieHandler() { + return Optional.empty(); + } + + @Override + public Optional executor() { + return Optional.of(exec); + } + + @Override + public Redirect followRedirects() { + return null; + } + + @Override + public Optional proxy() { + return Optional.empty(); + } + + @Override + public SSLContext sslContext() { + return null; + } + + @Override + public SSLParameters sslParameters() { + return null; + } + + @Override + public Version version() { + return Version.HTTP_1_1; + } + + @Override + public boolean awaitTermination(Duration duration) throws InterruptedException { + return exec.awaitTermination(duration.toMillis(), TimeUnit.MILLISECONDS); + } + + @Override + public HttpResponse send(java.net.http.HttpRequest request, BodyHandler responseBodyHandler) throws IOException, + InterruptedException { + return sendAsync(request, responseBodyHandler).join(); + } + + @Override + public CompletableFuture> sendAsync( + java.net.http.HttpRequest request, + BodyHandler responseBodyHandler, + PushPromiseHandler pushPromiseHandler + ) { + return sendAsync(request, responseBodyHandler); + } + + @Override + public CompletableFuture> sendAsync(HttpRequest request, BodyHandler responseBodyHandler) { + Arrays.stream(listeners).forEach(l -> l.onExecute(request, responseBodyHandler)); + + final CompletableFuture> callback = new CompletableFuture<>(); + exec.submit(() -> { + try { + HttpResponse httpResponse = responseOrException(request); + callback.complete((HttpResponse) httpResponse); + return (T) httpResponse; + } catch (Exception e) { + callback.completeExceptionally(e); + return null; + } + }); + + return callback; + } + }; + + return httpClient; + } + + private static HttpResponse responseOrException(HttpRequest request) throws Exception { + final HttpHost httpHost = new HttpHost(request.uri().getScheme(), request.uri().getHost(), request.uri().getPort()); + // return the desired status code or exception depending on the path + switch (request.uri().getPath()) { + case "/soe": + throw new SocketTimeoutException(httpHost.toString()); + case "/coe": + throw new HttpTimeoutException(httpHost.toString()); + case "/ioe": + throw new IOException(httpHost.toString()); + case "/closed": + throw new ClosedChannelException(); + case "/handshake": + throw new SSLHandshakeException(""); + case "/uri": + throw new URISyntaxException("", ""); + case "/runtime": + throw new RuntimeException(); + default: + int statusCode = Integer.parseInt(request.uri().getPath().substring(1)); + + // return the same body that was sent + final Object entity = BodyUtils.getBody(request).map(List::of).cache().block(); + final HttpResponse httpResponse = new HttpResponse<>() { + @Override + public int statusCode() { + return statusCode; + } + + @Override + public HttpRequest request() { + return request; + } + + @Override + public Optional> previousResponse() { + return Optional.empty(); + } + + @Override + public HttpHeaders headers() { + return request.headers(); + } + + @Override + public Object body() { + return entity; + } + + @Override + public Optional sslSession() { + return Optional.empty(); + } + + @Override + public URI uri() { + return request.uri(); + } + + @Override + public Version version() { + return request.version().orElse(Version.HTTP_1_1); + } + }; + + return httpResponse; + } + } + + /** + * Shutdown the executor so we don't leak threads into other test runs. + */ + @After + public void shutdownExec() { + exec.shutdown(); + } + + /** + * Verifies the content of the {@link HttpRequest} that's internally created and passed through to the http client + */ + @SuppressWarnings("unchecked") + public void testInternalHttpRequest() throws Exception { + int times = 0; + for (String httpMethod : getHttpMethods()) { + HttpRequest expectedRequest = performRandomRequest(httpMethod); + assertThat(requests.intValue(), equalTo(++times)); + + HttpRequest actualRequest = requestProducerCapture.get(); + assertEquals(expectedRequest.uri(), actualRequest.uri()); + assertEquals(expectedRequest.method(), actualRequest.method()); + assertEquals(expectedRequest.headers(), actualRequest.headers()); + + Object expectedEntity = BodyUtils.getBody(expectedRequest).block(); + if (expectedEntity != null) { + Object actualEntity = BodyUtils.getBody(actualRequest).block(); + assertEquals(expectedEntity, actualEntity); + } + } + } + + /** + * End to end test for ok status codes + */ + public void testOkStatusCodes() throws Exception { + for (String method : getHttpMethods()) { + for (int okStatusCode : getOkStatusCodes()) { + Response response = performRequestSyncOrAsync(restClient, Request.newRequest(method, "/" + okStatusCode).build()); + assertThat(response.statusLine().statusCode(), equalTo(okStatusCode)); + } + } + failureListener.assertNotCalled(); + } + + /** + * End to end test for error status codes: they should cause an exception to be thrown, apart from 404 with HEAD requests + */ + public void testErrorStatusCodes() throws Exception { + for (String method : getHttpMethods()) { + Set expectedIgnores = new HashSet<>(); + String ignoreParam = ""; + if ("HEAD".equalsIgnoreCase(method)) { + expectedIgnores.add(404); + } + if (randomBoolean()) { + int numIgnores = randomIntBetween(1, 3); + for (int i = 0; i < numIgnores; i++) { + Integer code = randomFrom(getAllErrorStatusCodes()); + expectedIgnores.add(code); + ignoreParam += code; + if (i < numIgnores - 1) { + ignoreParam += ","; + } + } + } + // error status codes should cause an exception to be thrown + for (int errorStatusCode : getAllErrorStatusCodes()) { + try { + Request.Builder builder = Request.newRequest(method, "/" + errorStatusCode); + if (false == ignoreParam.isEmpty()) { + builder = builder.withParameter("ignore", ignoreParam); + } + Response response = restClient.performRequest(builder.build()); + if (expectedIgnores.contains(errorStatusCode)) { + // no exception gets thrown although we got an error status code, as it was configured to be ignored + assertEquals(errorStatusCode, response.statusLine().statusCode()); + } else { + fail("request should have failed"); + } + } catch (ResponseException e) { + if (expectedIgnores.contains(errorStatusCode)) { + throw e; + } + assertEquals(errorStatusCode, e.getResponse().statusLine().statusCode()); + assertExceptionStackContainsCallingMethod(e); + } + if (errorStatusCode <= 500 || expectedIgnores.contains(errorStatusCode)) { + failureListener.assertNotCalled(); + } else { + failureListener.assertCalled(singletonList(node)); + } + } + } + } + + public void testPerformRequestIOExceptions() throws Exception { + for (String method : getHttpMethods()) { + // IOExceptions should be let bubble up + try { + restClient.performRequest(Request.newRequest(method, "/ioe").build()); + fail("request should have failed"); + } catch (IOException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/coe").build()); + fail("request should have failed"); + } catch (HttpTimeoutException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/soe").build()); + fail("request should have failed"); + } catch (SocketTimeoutException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/closed").build()); + fail("request should have failed"); + } catch (ClosedChannelException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + try { + restClient.performRequest(Request.newRequest(method, "/handshake").build()); + fail("request should have failed"); + } catch (SSLHandshakeException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + public void testPerformRequestRuntimeExceptions() throws Exception { + for (String method : getHttpMethods()) { + try { + restClient.performRequest(Request.newRequest(method, "/runtime").build()); + fail("request should have failed"); + } catch (RuntimeException e) { + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + public void testPerformRequestExceptions() throws Exception { + for (String method : getHttpMethods()) { + try { + restClient.performRequest(Request.newRequest(method, "/uri").build()); + fail("request should have failed"); + } catch (RuntimeException e) { + assertThat(e.getCause(), instanceOf(URISyntaxException.class)); + // And we do all that so the thrown exception has our method in the stacktrace + assertExceptionStackContainsCallingMethod(e); + } + failureListener.assertCalled(singletonList(node)); + } + } + + /** + * End to end test for request and response body. Exercises the mock http client ability to send back + * whatever body it has received. + */ + public void testBody() throws Exception { + String body = "{ \"field\": \"value\" }"; + BodyPublisher entity = BodyPublishers.ofString(body); + for (String method : Arrays.asList("DELETE", "GET", "PATCH", "POST", "PUT", "TRACE")) { + for (int okStatusCode : getOkStatusCodes()) { + Request request = Request.newRequest(method, "/" + okStatusCode).withEntity(entity).build(); + Response response = restClient.performRequest(request); + assertThat(response.statusLine().statusCode(), equalTo(okStatusCode)); + assertThat(BodyUtils.getBodyAsString(response), equalTo(body)); + } + for (int errorStatusCode : getAllErrorStatusCodes()) { + Request request = Request.newRequest(method, "/" + errorStatusCode).withEntity(entity).build(); + try { + restClient.performRequest(request); + fail("request should have failed"); + } catch (ResponseException e) { + Response response = e.getResponse(); + assertThat(response.statusLine().statusCode(), equalTo(errorStatusCode)); + assertThat(BodyUtils.getBodyAsString(response), equalTo(body)); + assertExceptionStackContainsCallingMethod(e); + } + } + } + } + + /** + * End to end test for request and response headers. Exercises the mock http client ability to send back + * whatever headers it has received. + */ + public void testHeaders() throws Exception { + for (String method : getHttpMethods()) { + final Map> requestHeaders = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + final int statusCode = randomStatusCode(getRandom()); + Request.Builder builder = Request.newRequest(method, "/" + statusCode); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (Map.Entry> requestHeader : requestHeaders.entrySet()) { + requestHeader.getValue().forEach(v -> options.addHeader(requestHeader.getKey(), v)); + } + builder = builder.withOptions(options); + Response esResponse; + try { + esResponse = performRequestSyncOrAsync(restClient, builder.build()); + } catch (ResponseException e) { + esResponse = e.getResponse(); + } + assertThat(esResponse.statusLine().statusCode(), equalTo(statusCode)); + assertHeaders(defaultHeaders, requestHeaders, esResponse.headers(), Collections.emptySet()); + assertFalse(esResponse.hasWarnings()); + } + } + + public void testDeprecationWarnings() throws Exception { + String chars = randomAsciiAlphanumOfLength(5); + assertDeprecationWarnings(singletonList("poorly formatted " + chars), singletonList("poorly formatted " + chars)); + assertDeprecationWarnings(singletonList(formatWarningWithoutDate(chars)), singletonList(chars)); + assertDeprecationWarnings(singletonList(formatWarning(chars)), singletonList(chars)); + assertDeprecationWarnings( + Arrays.asList(formatWarning(chars), "another one", "and another"), + Arrays.asList(chars, "another one", "and another") + ); + assertDeprecationWarnings(Arrays.asList("ignorable one", "and another"), Arrays.asList("ignorable one", "and another")); + assertDeprecationWarnings(singletonList("exact"), singletonList("exact")); + assertDeprecationWarnings(Collections.emptyList(), Collections.emptyList()); + + String proxyWarning = "112 - \"network down\" \"Sat, 25 Aug 2012 23:34:45 GMT\""; + assertDeprecationWarnings(singletonList(proxyWarning), singletonList(proxyWarning)); + } + + private enum DeprecationWarningOption { + PERMISSIVE { + protected WarningsHandler warningsHandler() { + return WarningsHandler.PERMISSIVE; + } + }, + STRICT { + protected WarningsHandler warningsHandler() { + return WarningsHandler.STRICT; + } + }, + FILTERED { + protected WarningsHandler warningsHandler() { + return new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + for (String warning : warnings) { + if (false == warning.startsWith("ignorable")) { + return true; + } + } + return false; + } + }; + } + }, + EXACT { + protected WarningsHandler warningsHandler() { + return new WarningsHandler() { + @Override + public boolean warningsShouldFailRequest(List warnings) { + return false == warnings.equals(Arrays.asList("exact")); + } + }; + } + }; + + protected abstract WarningsHandler warningsHandler(); + } + + private void assertDeprecationWarnings(List warningHeaderTexts, List warningBodyTexts) throws Exception { + String method = randomFrom(getHttpMethods()); + Request.Builder builder = Request.newRequest(method, "/200"); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + for (String warningHeaderText : warningHeaderTexts) { + options.addHeader("Warning", warningHeaderText); + } + + final boolean expectFailure; + if (randomBoolean()) { + logger.info("checking strictWarningsMode=[" + strictDeprecationMode + "] and warnings=" + warningBodyTexts); + expectFailure = strictDeprecationMode && false == warningBodyTexts.isEmpty(); + } else { + DeprecationWarningOption warningOption = randomFrom(DeprecationWarningOption.values()); + logger.info("checking warningOption=" + warningOption + " and warnings=" + warningBodyTexts); + options.setWarningsHandler(warningOption.warningsHandler()); + expectFailure = warningOption.warningsHandler().warningsShouldFailRequest(warningBodyTexts); + } + builder = builder.withOptions(options); + + Response response; + if (expectFailure) { + try { + performRequestSyncOrAsync(restClient, builder.build()); + fail("expected WarningFailureException from warnings"); + return; + } catch (WarningFailureException e) { + if (false == warningBodyTexts.isEmpty()) { + assertThat(e.getMessage(), containsString("\nWarnings: " + warningBodyTexts)); + } + response = e.getResponse(); + } + } else { + response = performRequestSyncOrAsync(restClient, builder.build()); + } + assertEquals(false == warningBodyTexts.isEmpty(), response.hasWarnings()); + assertEquals(warningBodyTexts, response.warnings()); + } + + /** + * Emulates OpenSearch's HeaderWarningLogger.formatWarning in simple + * cases. We don't have that available because we're testing against 1.7. + */ + private static String formatWarningWithoutDate(String warningBody) { + final String hash = new String(new byte[40], StandardCharsets.UTF_8).replace('\0', 'e'); + return "299 OpenSearch-1.2.2-SNAPSHOT-" + hash + " \"" + warningBody + "\""; + } + + private static String formatWarning(String warningBody) { + return formatWarningWithoutDate(warningBody) + " \"Mon, 01 Jan 2001 00:00:00 GMT\""; + } + + private HttpRequest performRandomRequest(String method) throws Exception { + String uriAsString = "/" + randomStatusCode(getRandom()); + Request.Builder builder = Request.newRequest(method, uriAsString); + + Map params = new HashMap<>(); + if (randomBoolean()) { + int numParams = randomIntBetween(1, 3); + for (int i = 0; i < numParams; i++) { + String name = "param-" + i; + String value = randomAsciiAlphanumOfLengthBetween(3, 10); + builder = builder.withParameter(name, value); + params.put(name, value); + } + } + if (randomBoolean()) { + // randomly add some ignore parameter, which doesn't get sent as part of the request + String ignore = Integer.toString(randomFrom(RestClientTestUtil.getAllErrorStatusCodes())); + if (randomBoolean()) { + ignore += "," + randomFrom(RestClientTestUtil.getAllErrorStatusCodes()); + } + builder = builder.withParameter("ignore", ignore); + } + + final String additionalQuery = params.entrySet() + .stream() + .map(e -> e.getKey() + "=" + URLEncoder.encode(e.getValue(), StandardCharsets.UTF_8)) + .collect(Collectors.joining("&")); + + URI uri = new URI("http", null, "localhost", 9200, uriAsString, additionalQuery, null); + BodyPublisher bodyPublisher = BodyPublishers.noBody(); + if (getRandom().nextBoolean()) { + bodyPublisher = BodyPublishers.ofString(randomAsciiAlphanumOfLengthBetween(10, 100)); + builder = builder.withEntity(bodyPublisher); + } + + HttpRequest.Builder expectedRequest = HttpRequest.newBuilder(uri).method(method, bodyPublisher); + final Set uniqueNames = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + if (randomBoolean()) { + Map> headers = RestClientTestUtil.randomHeaders(getRandom(), "Header"); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + options.addHeaders(headers); + for (Map.Entry> header : headers.entrySet()) { + header.getValue().forEach(v -> expectedRequest.header(header.getKey(), v)); + uniqueNames.add(header.getKey()); + } + builder = builder.withOptions(options); + } + for (Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + // request level headers override default headers + if (uniqueNames.contains(defaultHeader.getKey()) == false) { + defaultHeader.getValue().forEach(v -> expectedRequest.header(defaultHeader.getKey(), v)); + } + } + try { + performRequestSyncOrAsync(restClient, builder.build()); + } catch (Exception e) { + // all good + } + return expectedRequest.build(); + } + + static Response performRequestSyncOrAsync(RestHttpClient restClient, Request request) throws Exception { + // randomize between sync and async methods + if (randomBoolean()) { + return restClient.performRequest(request); + } else { + final AtomicReference exceptionRef = new AtomicReference<>(); + final AtomicReference responseRef = new AtomicReference<>(); + final CountDownLatch latch = new CountDownLatch(1); + restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + responseRef.set(response); + latch.countDown(); + + } + + @Override + public void onFailure(Exception exception) { + exceptionRef.set(exception); + latch.countDown(); + } + }); + latch.await(); + if (exceptionRef.get() != null) { + throw exceptionRef.get(); + } + return responseRef.get(); + } + } + + /** + * Asserts that the provided {@linkplain Exception} contains the method + * that called this somewhere on its stack. This is + * normally the case for synchronous calls but {@link RestHttpClient} performs + * synchronous calls by performing asynchronous calls and blocking the + * current thread until the call returns so it has to take special care + * to make sure that the caller shows up in the exception. We use this + * assertion to make sure that we don't break that "special care". + */ + private static void assertExceptionStackContainsCallingMethod(Throwable t) { + // 0 is getStackTrace + // 1 is this method + // 2 is the caller, what we want + StackTraceElement myMethod = Thread.currentThread().getStackTrace()[2]; + for (StackTraceElement se : t.getStackTrace()) { + if (se.getClassName().equals(myMethod.getClassName()) && se.getMethodName().equals(myMethod.getMethodName())) { + return; + } + } + StringWriter stack = new StringWriter(); + t.printStackTrace(new PrintWriter(stack)); + fail("didn't find the calling method (looks like " + myMethod + ") in:\n" + stack); + } +} diff --git a/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java new file mode 100644 index 0000000000000..e7ddb131b4205 --- /dev/null +++ b/client/rest-http-client/src/test/java/org/opensearch/internal/httpclient/RestHttpClientTestCase.java @@ -0,0 +1,99 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.internal.httpclient; + +import com.carrotsearch.randomizedtesting.JUnit3MethodProvider; +import com.carrotsearch.randomizedtesting.MixWithSuiteName; +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.SeedDecorators; +import com.carrotsearch.randomizedtesting.annotations.TestMethodProviders; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakAction; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakGroup; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakLingering; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakZombies; +import com.carrotsearch.randomizedtesting.annotations.TimeoutSuite; + +import java.net.http.HttpHeaders; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +@TestMethodProviders({ JUnit3MethodProvider.class }) +@SeedDecorators({ MixWithSuiteName.class }) // See LUCENE-3995 for rationale. +@ThreadLeakScope(ThreadLeakScope.Scope.SUITE) +@ThreadLeakGroup(ThreadLeakGroup.Group.MAIN) +@ThreadLeakAction({ ThreadLeakAction.Action.WARN, ThreadLeakAction.Action.INTERRUPT }) +@ThreadLeakZombies(ThreadLeakZombies.Consequence.IGNORE_REMAINING_TESTS) +@ThreadLeakLingering(linger = 5000) // 5 sec lingering +@ThreadLeakFilters(filters = { HttpClientThreadLeakFilter.class, BouncyCastleThreadFilter.class }) +@TimeoutSuite(millis = 2 * 60 * 60 * 1000) +public abstract class RestHttpClientTestCase extends RandomizedTest { + /** + * Assert that the actual headers are the expected ones given the original default and request headers. Some headers can be ignored, + * for instance in case the http client is adding its own automatically. + * + * @param defaultHeaders the default headers set to the REST client instance + * @param requestHeaders the request headers sent with a particular request + * @param actualHeaders the actual headers as a result of the provided default and request headers + * @param ignoreHeaders header keys to be ignored as they are not part of default nor request headers, yet they + * will be part of the actual ones + */ + protected static void assertHeaders( + final Map> defaultHeaders, + final Map> requestHeaders, + final HttpHeaders actualHeaders, + final Set ignoreHeaders + ) { + final Map> expectedHeaders = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); + final Set requestHeaderKeys = new HashSet<>(); + for (final Map.Entry> header : requestHeaders.entrySet()) { + final String name = header.getKey(); + addValueToListEntry(expectedHeaders, name, header.getValue()); + requestHeaderKeys.add(name); + } + for (final Map.Entry> defaultHeader : defaultHeaders.entrySet()) { + final String name = defaultHeader.getKey(); + if (requestHeaderKeys.contains(name) == false) { + addValueToListEntry(expectedHeaders, name, defaultHeader.getValue()); + } + } + Set actualIgnoredHeaders = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + for (final Map.Entry> responseHeader : actualHeaders.map().entrySet()) { + final String name = responseHeader.getKey(); + if (ignoreHeaders.contains(name)) { + expectedHeaders.remove(name); + actualIgnoredHeaders.add(name); + continue; + } + final String value = responseHeader.getValue().getFirst(); + final List values = expectedHeaders.get(name); + assertNotNull("found response header [" + name + "] that wasn't originally sent: " + value, values); + assertTrue("found incorrect response header [" + name + "]: " + value, values.remove(value)); + if (values.isEmpty()) { + expectedHeaders.remove(name); + } + } + assertEquals("some headers meant to be ignored were not part of the actual headers", ignoreHeaders, actualIgnoredHeaders); + assertTrue("some headers that were sent weren't returned " + expectedHeaders, expectedHeaders.isEmpty()); + } + + private static void addValueToListEntry(final Map> map, final String name, final List values) { + map.computeIfAbsent(name, k -> new ArrayList<>()).addAll(values); + } +} diff --git a/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java b/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java index fc73e03201e84..92b7e740da0e1 100644 --- a/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java +++ b/client/rest/src/test/java/org/opensearch/client/RestClientSingleHostIntegTests.java @@ -483,6 +483,18 @@ public void testEncodeParams() throws Exception { Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); assertEquals(pathPrefix + "/200?routing=foo%5Ebar", response.getRequestLine().getUri()); } + { + Request request = new Request("PUT", "/200"); + request.addParameter("pretty", null); + Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty", response.getRequestLine().getUri()); + } + { + Request request = new Request("PUT", "/200"); + request.addParameter("pretty", ""); + Response response = RestClientSingleHostTests.performRequestSyncOrAsync(restClient, request); + assertEquals(pathPrefix + "/200?pretty=", response.getRequestLine().getUri()); + } } /** diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 483412f03b386..875d0bca23966 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -97,7 +97,7 @@ jzlib = "1.1.3" resteasy = "6.2.4.Final" # opentelemetry dependencies -opentelemetry = "1.62.0" +opentelemetry = "1.63.0" opentelemetrysemconv = "1.41.0" # arrow dependencies diff --git a/libs/arrow-spi/build.gradle b/libs/arrow-spi/build.gradle index abf7eecf84c77..c1a716c6f72c4 100644 --- a/libs/arrow-spi/build.gradle +++ b/libs/arrow-spi/build.gradle @@ -11,7 +11,9 @@ apply plugin: 'opensearch.publish' dependencies { api project(':libs:opensearch-core') api project(':libs:opensearch-common') - testImplementation project(':test:framework') + testImplementation(project(':test:framework')) { + exclude group: 'org.opensearch', module: 'opensearch-arrow-spi' + } } tasks.named('forbiddenApisMain').configure { diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java index 89d6866da2c89..874d016cc40bb 100644 --- a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocator.java @@ -9,36 +9,35 @@ package org.opensearch.arrow.spi; import java.io.Closeable; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Supplier; /** - * Arrow-agnostic interface for a hierarchical native memory allocator. + * Unified native memory allocator interface. * - *

The implementation (backed by Arrow's {@code RootAllocator}) is provided by - * a plugin. The SPI allows other subsystems to interact with the allocator - * without depending on Arrow classes. - * - *

Plugins that need Arrow allocators obtain the implementation via - * service lookup or plugin extension and call {@link #getOrCreatePool} to - * register their pool. + *

Manages memory pools under a shared budget. Each pool has a minimum + * guaranteed allocation and a maximum burst limit. Implementations may + * redistribute unused capacity across pools. * * @opensearch.api */ public interface NativeAllocator extends Closeable { /** - * Returns the named pool, creating it on first access with the given limit. - * Subsequent calls with the same name return the same pool (first-call limit wins). + * Returns the named pool, creating it on first access. + * Subsequent calls with the same name return the existing pool (first-call config wins). * - * @param poolName logical pool name (e.g., "query", "flight") - * @param limit maximum bytes this pool can allocate in aggregate + * @param poolName logical pool name + * @param min minimum guaranteed bytes + * @param max maximum bytes this pool can allocate + * @param group the group this pool belongs to for aggregated stats, or null * @return an opaque pool handle */ - PoolHandle getOrCreatePool(String poolName, long limit); + PoolHandle getOrCreatePool(String poolName, long min, long max, PoolGroup group); /** - * Updates the limit of an existing pool. Children of the pool allocator - * inherit the change automatically via Arrow's parent-cap check at - * allocation time — no notification SPI is needed. + * Updates the effective limit of an existing pool. * * @param poolName logical pool name * @param newLimit new maximum bytes for the pool @@ -46,15 +45,78 @@ public interface NativeAllocator extends Closeable { void setPoolLimit(String poolName, long newLimit); /** - * Sets the root-level memory limit for the entire allocator. + * Registers a virtual pool with initial min/max and a callback + * invoked when the pool's limit changes. + * + * @param poolName logical pool name + * @param min minimum guaranteed bytes + * @param max initial maximum bytes (the pool's starting limit) + * @param group the group this pool belongs to for aggregated stats + * @param limitSetter callback invoked when the pool limit changes + * @return a handle to update stats from the native layer + */ + VirtualPoolHandle registerVirtualPool(String poolName, long min, long max, PoolGroup group, Consumer limitSetter); + + /** + * Updates the minimum guaranteed bytes for a pool. + * + * @param poolName logical pool name + * @param newMin new minimum bytes + */ + void setPoolMin(String poolName, long newMin); + + /** + * Returns all registered pool names. + */ + Set getAllPoolNames(); + + /** + * Adds a callback invoked before stats collection to refresh pool usage data. + * + * @param refresher runnable that updates pool stats + */ + void addStatsRefresher(Runnable refresher); + + /** + * Sets the supplier for process-wide native memory stats. + * + * @param supplier returns [allocatedBytes, residentBytes] + */ + void setNativeMemoryStatsSupplier(Supplier supplier); + + /** + * Registers a listener invoked when the effective limit of a pool group changes. + * The consumer receives the new total effective limit (sum of all pools in the group). * - * @param limit new maximum bytes for the root allocator + * @param group the pool group to listen on + * @param listener consumer that receives the new grouped limit in bytes + */ + void addPoolGroupLimitListener(PoolGroup group, Consumer listener); + + /** + * Handle for a virtual pool. Plugins update stats via this handle. */ - void setRootLimit(long limit); + interface VirtualPoolHandle { + /** + * Update the current usage stats. + * + * @param allocatedBytes current allocated bytes + * @param peakBytes peak allocated bytes + */ + void updateStats(long allocatedBytes, long peakBytes); + + /** Returns current allocated bytes. */ + long allocatedBytes(); + + /** Returns peak allocated bytes. */ + long peakBytes(); + + /** Returns current limit. */ + long limit(); + } /** - * Opaque handle to a memory pool. Plugins downcast to the concrete type - * (e.g., Arrow's {@code BufferAllocator}) in the implementation layer. + * Opaque handle to a memory pool. */ interface PoolHandle { @@ -63,28 +125,20 @@ interface PoolHandle { * * @param childName name for debugging * @param childLimit maximum bytes for the child - * @return an opaque child handle (downcast to BufferAllocator in Arrow impl) + * @return a child handle */ PoolHandle newChild(String childName, long childLimit); - /** - * Returns the current allocated bytes for this pool/child. - */ + /** Returns the current allocated bytes. */ long allocatedBytes(); - /** - * Returns the peak memory allocation. - */ + /** Returns the peak memory allocation. */ long peakBytes(); - /** - * Returns the configured limit. - */ + /** Returns the configured limit. */ long limit(); - /** - * Releases this allocation handle. - */ + /** Releases this allocation handle. */ void close(); } } diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java index 98f991cb86704..29dba48e9f165 100644 --- a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfig.java @@ -33,9 +33,6 @@ public final class NativeAllocatorPoolConfig { /** Pool name for query-execution memory (analytics-engine fragments and per-query allocators). */ public static final String POOL_QUERY = "query"; - /** Setting key for the root allocator limit. */ - public static final String SETTING_ROOT_LIMIT = "native.allocator.root.limit"; - /** Setting key for the Flight pool minimum. */ public static final String SETTING_FLIGHT_MIN = "native.allocator.pool.flight.min"; /** Setting key for the Flight pool maximum. */ diff --git a/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java new file mode 100644 index 0000000000000..d292b42ae595f --- /dev/null +++ b/libs/arrow-spi/src/main/java/org/opensearch/arrow/spi/PoolGroup.java @@ -0,0 +1,37 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.spi; + +/** + * Groups that memory pools belong to for aggregated customer-facing stats. + * Each pool is assigned to exactly one group at registration time. + * + * @opensearch.api + */ +public enum PoolGroup { + /** Arrow Flight transport pool group. */ + TRANSPORT("transport"), + /** Query and analytics execution pool group. */ + SEARCH("search"), + /** Ingest and write path pool group. */ + INDEXING("indexing"), + /** Background merge operations pool group. */ + MERGE("merge"); + + private final String name; + + PoolGroup(String name) { + this.name = name; + } + + /** Returns the group name used in stats output. */ + public String getName() { + return name; + } +} diff --git a/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java b/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java index a21ca8ff54943..025b41a603212 100644 --- a/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java +++ b/libs/arrow-spi/src/test/java/org/opensearch/arrow/spi/NativeAllocatorPoolConfigTests.java @@ -26,8 +26,4 @@ public void testSettingKeys() { assertEquals("native.allocator.pool.query.min", NativeAllocatorPoolConfig.SETTING_QUERY_MIN); assertEquals("native.allocator.pool.query.max", NativeAllocatorPoolConfig.SETTING_QUERY_MAX); } - - public void testRootSettingKey() { - assertEquals("native.allocator.root.limit", NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT); - } } diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java index a9bd9968b5884..6a2c52a46f9ce 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBasePlugin.java @@ -9,11 +9,16 @@ package org.opensearch.arrow.allocator; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.IndexScopedSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.common.settings.SettingsFilter; +import org.opensearch.common.util.concurrent.FutureUtils; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.core.xcontent.NamedXContentRegistry; @@ -22,10 +27,14 @@ import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.plugin.stats.NativeAllocatorStatsRegistry; +import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.ExtensiblePlugin; import org.opensearch.plugins.Plugin; import org.opensearch.repositories.RepositoriesService; +import org.opensearch.rest.RestController; +import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; +import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.Client; import org.opensearch.watcher.ResourceWatcherService; @@ -33,207 +42,133 @@ import java.io.IOException; import java.util.Collection; import java.util.List; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; import java.util.function.Supplier; /** * Top-level plugin that owns the unified Arrow-backed native memory allocator. * - *

All Arrow-consuming plugins (arrow-flight-rpc, parquet-data-format) extend - * this plugin to share one {@link ArrowNativeAllocator} and its classloader. - * - *

Each pool has a min (guaranteed floor) and max (burst ceiling). The rebalancer - * ensures every pool can always allocate up to its min, and distributes unused - * capacity allowing pools to grow up to their max. + *

All Arrow-consuming plugins extend this plugin to share one + * {@link ArrowNativeAllocator} and its classloader. */ -public class ArrowBasePlugin extends Plugin implements ExtensiblePlugin { +public class ArrowBasePlugin extends Plugin implements ExtensiblePlugin, ActionPlugin { /** Creates the plugin. */ public ArrowBasePlugin() {} - /** - * Maximum bytes for the root Arrow allocator. - * - *

When unset, the default is 20% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #deriveRootLimitDefault}. The Arrow framework gets a small fraction of the - * native budget because the dominant consumer of native memory in analytics workloads - * is the DataFusion Rust runtime (~75% of {@code node.native_memory.limit}), not Arrow. - * If AC is unconfigured (limit = 0), the default is {@link Long#MAX_VALUE}, preserving - * pre-AC behaviour. - */ - public static final Setting ROOT_LIMIT_SETTING = new Setting<>( - NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, - ArrowBasePlugin::deriveRootLimitDefault, - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT + "] must be >= 0, got " + v); - } - return v; - }, + // ─── Settings ──────────────────────────────────────────────────────────────── + + /** Whether the NativeMemoryRebalancer is enabled. */ + public static final Setting REBALANCER_ENABLED_SETTING = Setting.boolSetting( + "native.allocator.rebalancer.enabled", + true, Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Computes the default for {@link #ROOT_LIMIT_SETTING} as 20% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}. The Arrow framework's - * hard cap covers only Arrow allocations — DataFusion's Rust runtime is a sibling of - * Arrow root and gets the larger share of the native budget (see - * {@code DataFusionPlugin#deriveMemoryPoolLimitDefault}). - * - *

Returns the bytes-as-string representation expected by the {@link Setting} parser. - * If the AC limit is unset (== 0), the default is {@link Long#MAX_VALUE} — unbounded — - * preserving pre-AC behaviour. - */ - static String deriveRootLimitDefault(Settings settings) { - ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); - if (nativeLimit.getBytes() <= 0) { - return Long.toString(Long.MAX_VALUE); - } - return Long.toString(nativeLimit.getBytes() * 20 / 100); - } + /** Interval in seconds between pool rebalance cycles. 0 disables rebalancing. */ + public static final Setting REBALANCE_INTERVAL_SETTING = Setting.longSetting( + "native.allocator.rebalance.interval_seconds", + 5L, + 0L, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Pool utilization above this triggers growth. */ + public static final Setting PRESSURE_THRESHOLD_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.pressure_threshold", + 0.75, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Pool utilization below this means pool can give back capacity. */ + public static final Setting IDLE_THRESHOLD_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.idle_threshold", + 0.50, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); - /** Minimum guaranteed bytes for the Flight pool. */ - public static final Setting FLIGHT_MIN_SETTING = Setting.longSetting( + /** Factor to shrink idle pools by (new limit = limit * (1 - shrink_factor)). */ + public static final Setting SHRINK_FACTOR_SETTING = Setting.doubleSetting( + "native.allocator.rebalancer.shrink_factor", + 0.10, + 0.0, + 1.0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** Minimum guaranteed bytes for the Flight pool. Default is 2% of budget. */ + public static final Setting FLIGHT_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 2), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the Flight pool can burst to. Default is 5% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Matches the partitioning model documented in PR #21732. - */ + /** Maximum bytes the Flight pool can burst to. Default is 5% of budget. */ public static final Setting FLIGHT_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, s -> derivePoolMaxDefault(s, 5), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** Minimum guaranteed bytes for the ingest pool. */ - public static final Setting INGEST_MIN_SETTING = Setting.longSetting( + /** Minimum guaranteed bytes for the ingest pool. Default is 4% of budget. */ + public static final Setting INGEST_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_INGEST_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 4), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_INGEST_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the ingest pool can burst to. Default is 8% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Ingest gets a larger fraction than Flight/Query because parquet VSR - * allocators dominate write-path memory usage — see partitioning model in PR #21732. - */ + /** Maximum bytes the ingest pool can burst to. Default is 8% of budget. */ public static final Setting INGEST_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_INGEST_MAX, s -> derivePoolMaxDefault(s, 8), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_INGEST_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_INGEST_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Minimum guaranteed bytes for the query pool. Honored by the rebalancer (when - * enabled) — sets a floor below which the rebalancer will not shrink the pool. - * Has no effect when rebalancing is disabled. - */ - public static final Setting QUERY_MIN_SETTING = Setting.longSetting( + /** Minimum guaranteed bytes for the query pool. Default is 2% of budget. */ + public static final Setting QUERY_MIN_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_QUERY_MIN, - 0L, - 0L, + s -> derivePoolMinDefault(s, 2), + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MIN), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Maximum bytes the query pool can allocate. Default is 5% of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING}; see - * {@link #derivePoolMaxDefault}. Falls back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Enforced by Arrow's child-allocator limit — analytics-engine's - * per-query allocators are children of this pool, so the sum of in-flight per-query - * allocations is capped here. - * - *

Note: each individual analytics query is also bounded by - * {@code analytics.exec.QueryContext} per-query limit (currently the constant - * {@code DEFAULT_PER_QUERY_MEMORY_LIMIT = 256 MB}). Lowering {@code QUERY_MAX} - * below {@code 256 MB × concurrent-queries} can starve queries even when each - * individual query is within its per-query limit. - */ + /** Maximum bytes the query pool can allocate. Default is 5% of budget. */ public static final Setting QUERY_MAX_SETTING = new Setting<>( NativeAllocatorPoolConfig.SETTING_QUERY_MAX, s -> derivePoolMaxDefault(s, 5), - s -> { - long v = Long.parseLong(s); - if (v < 0) { - throw new IllegalArgumentException("Setting [" + NativeAllocatorPoolConfig.SETTING_QUERY_MAX + "] must be >= 0, got " + v); - } - return v; - }, + s -> parseNonNegativeLong(s, NativeAllocatorPoolConfig.SETTING_QUERY_MAX), Setting.Property.NodeScope, Setting.Property.Dynamic ); - /** - * Computes the default for a pool max as a percentage of - * {@link ResourceTrackerSettings#NODE_NATIVE_MEMORY_LIMIT_SETTING} (the operator's - * declared off-heap budget), falling back to {@link Long#MAX_VALUE} when AC is - * unconfigured. Returns the bytes-as-string representation expected by the - * {@link Setting} parser. - * - *

Pools are anchored to {@code node.native_memory.limit} rather than to - * {@link #ROOT_LIMIT_SETTING} so the diagrammed partitioning (PR #21732) holds: - * sum of pool maxes (5+8+5 = 18% of native_memory.limit) fits within the framework - * root cap (20% of native_memory.limit) by default. Operator overrides of - * {@code root.limit} that drop it below {@code sum(pool.max)} are caught by the - * grouped validator. - * - *

The fraction is taken straight from {@code node.native_memory.limit}, not from - * {@code limit - buffer_percent}. {@code buffer_percent} is an admission-control - * throttle margin, not a framework budget reduction. - * - * @param settings node settings - * @param percent fraction of {@code node.native_memory.limit} the pool max defaults to - */ - static String derivePoolMaxDefault(Settings settings, int percent) { - ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); - if (nativeLimit.getBytes() <= 0) { - return Long.toString(Long.MAX_VALUE); - } - long pool = Math.max(0L, nativeLimit.getBytes() * percent / 100); - return Long.toString(pool); - } - - /** Interval in seconds between pool rebalance cycles. 0 disables rebalancing. */ - public static final Setting REBALANCE_INTERVAL_SETTING = Setting.longSetting( - "native.allocator.rebalance.interval_seconds", - 0L, - 0L, - Setting.Property.NodeScope, - Setting.Property.Dynamic - ); + // ─── Instance state ────────────────────────────────────────────────────────── private volatile ArrowNativeAllocator allocator; + private volatile ScheduledExecutorService rebalancerScheduler; + private volatile ScheduledFuture rebalanceTask; + private volatile NativeMemoryRebalancer rebalancer; + + // ─── Plugin lifecycle ──────────────────────────────────────────────────────── @Override public Collection createComponents( @@ -251,12 +186,11 @@ public Collection createComponents( ) { Settings settings = environment.settings(); ClusterSettings cs = clusterService.getClusterSettings(); - ArrowNativeAllocator built = buildAllocator(settings, cs); + Supplier budgetSupplier = () -> ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(clusterService.getSettings()) + .getBytes(); + ArrowNativeAllocator built = buildAllocator(settings, cs, budgetSupplier); this.allocator = built; - // Publish a NativeAllocatorStatsRegistry alongside the allocator so the server-side - // NodeService can discover the supplier via pluginComponents (instanceof filter) without - // taking a compile-time dependency on this plugin. The lambda re-reads `this.allocator` - // each invocation, so after close() nulls the field, the supplier returns null cleanly. + Supplier statsSupplier = () -> { ArrowNativeAllocator a = this.allocator; return a != null ? a.stats() : null; @@ -264,96 +198,177 @@ public Collection createComponents( return List.of(built, new NativeAllocatorStatsRegistry(statsSupplier)); } + @Override + public List> getSettings() { + return List.of( + FLIGHT_MIN_SETTING, + FLIGHT_MAX_SETTING, + INGEST_MIN_SETTING, + INGEST_MAX_SETTING, + QUERY_MIN_SETTING, + QUERY_MAX_SETTING, + REBALANCE_INTERVAL_SETTING, + REBALANCER_ENABLED_SETTING, + PRESSURE_THRESHOLD_SETTING, + IDLE_THRESHOLD_SETTING, + SHRINK_FACTOR_SETTING + ); + } + + @Override + public List getRestHandlers( + Settings settings, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster + ) { + Supplier statsSupplier = () -> allocator != null ? allocator.stats() : null; + return List.of(new ArrowBaseStatsAction(statsSupplier)); + } + + @Override + public void close() throws IOException { + if (rebalancerScheduler != null) { + rebalancerScheduler.shutdownNow(); + } + if (allocator != null) { + allocator.close(); + allocator = null; + } + } + + // ─── Package-private (visible for tests) ───────────────────────────────────── + /** - * Constructs the allocator and wires its pools and dynamic-update consumers from - * a pure {@code (Settings, ClusterSettings)} pair. Package-private so unit tests - * can exercise the full wiring without a heavyweight {@link ClusterService} - * fixture — mirrors the shape of {@link #registerSettingsUpdateConsumers} which - * is already test-friendly for the same reason. + * Constructs the allocator and wires its pools and the rebalancer. */ - static ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs) { - long rootLimit = ROOT_LIMIT_SETTING.get(settings); - ArrowNativeAllocator allocator = new ArrowNativeAllocator(rootLimit); - allocator.setRebalanceInterval(REBALANCE_INTERVAL_SETTING.get(settings)); + ArrowNativeAllocator buildAllocator(Settings settings, ClusterSettings cs, Supplier budgetSupplier) { + ArrowNativeAllocator allocator = new ArrowNativeAllocator(); - // Single source of truth for cross-setting invariants — same logic runs on - // dynamic updates via the grouped consumer below. - validateUpdate(settings); + // Set budget for validation + long nativeBudget = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings).getBytes(); + if (nativeBudget > 0) { + allocator.setBudget(nativeBudget); + } + // Validate min < max for each pool + validateMinMax(NativeAllocatorPoolConfig.POOL_FLIGHT, FLIGHT_MIN_SETTING.get(settings), FLIGHT_MAX_SETTING.get(settings)); + validateMinMax(NativeAllocatorPoolConfig.POOL_INGEST, INGEST_MIN_SETTING.get(settings), INGEST_MAX_SETTING.get(settings)); + validateMinMax(NativeAllocatorPoolConfig.POOL_QUERY, QUERY_MIN_SETTING.get(settings), QUERY_MAX_SETTING.get(settings)); + + // Create pools (always start at max) allocator.getOrCreatePool( NativeAllocatorPoolConfig.POOL_FLIGHT, FLIGHT_MIN_SETTING.get(settings), - FLIGHT_MAX_SETTING.get(settings) + FLIGHT_MAX_SETTING.get(settings), + PoolGroup.TRANSPORT ); allocator.getOrCreatePool( NativeAllocatorPoolConfig.POOL_INGEST, INGEST_MIN_SETTING.get(settings), - INGEST_MAX_SETTING.get(settings) + INGEST_MAX_SETTING.get(settings), + PoolGroup.INDEXING + ); + allocator.getOrCreatePool( + NativeAllocatorPoolConfig.POOL_QUERY, + QUERY_MIN_SETTING.get(settings), + QUERY_MAX_SETTING.get(settings), + PoolGroup.SEARCH ); - allocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_QUERY, QUERY_MIN_SETTING.get(settings), QUERY_MAX_SETTING.get(settings)); - registerSettingsUpdateConsumers(cs, allocator); + // Register dynamic setting consumers for min/max changes + cs.addSettingsUpdateConsumer(FLIGHT_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT, newMin)); + cs.addSettingsUpdateConsumer(FLIGHT_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_FLIGHT, newMax)); + cs.addSettingsUpdateConsumer(INGEST_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_INGEST, newMin)); + cs.addSettingsUpdateConsumer(INGEST_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, newMax)); + cs.addSettingsUpdateConsumer(QUERY_MIN_SETTING, newMin -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_QUERY, newMin)); + cs.addSettingsUpdateConsumer(QUERY_MAX_SETTING, newMax -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, newMax)); + + // Register dynamic consumer for rebalancer enable/disable + cs.addSettingsUpdateConsumer(REBALANCER_ENABLED_SETTING, enabled -> { + if (enabled == false) { + cancelRebalanceTask(); + allocator.resetAllPoolsToMax(); + } else { + startRebalancer(allocator, budgetSupplier, cs.get(REBALANCE_INTERVAL_SETTING)); + } + }); + + // Set up the rebalancer if enabled + if (REBALANCER_ENABLED_SETTING.get(settings)) { + startRebalancer(allocator, budgetSupplier, REBALANCE_INTERVAL_SETTING.get(settings)); + } + + // Register dynamic consumer for interval changes + cs.addSettingsUpdateConsumer(REBALANCE_INTERVAL_SETTING, this::updateRebalanceInterval); + + // Register dynamic consumers for threshold changes + cs.addSettingsUpdateConsumer(PRESSURE_THRESHOLD_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setPressureThreshold(value); + }); + cs.addSettingsUpdateConsumer(IDLE_THRESHOLD_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setIdleThreshold(value); + }); + cs.addSettingsUpdateConsumer(SHRINK_FACTOR_SETTING, value -> { + NativeMemoryRebalancer r = this.rebalancer; + if (r != null) r.setShrinkFactor(value); + }); + return allocator; } - /** - * Registers cluster-settings update consumers that propagate dynamic setting changes - * into the live {@link ArrowNativeAllocator}. Package-private so unit tests can exercise - * the wiring with a real {@link ClusterSettings} instance — the test that asserts a PUT - * lands on the allocator is what catches a future regression where one of these lines - * is accidentally removed. - */ - static void registerSettingsUpdateConsumers(ClusterSettings cs, ArrowNativeAllocator allocator) { - cs.addSettingsUpdateConsumer(ROOT_LIMIT_SETTING, allocator::setRootLimit); - cs.addSettingsUpdateConsumer(REBALANCE_INTERVAL_SETTING, allocator::setRebalanceInterval); - cs.addSettingsUpdateConsumer(FLIGHT_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_FLIGHT, v)); - cs.addSettingsUpdateConsumer(FLIGHT_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT, v)); - cs.addSettingsUpdateConsumer(INGEST_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_INGEST, v)); - cs.addSettingsUpdateConsumer(INGEST_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_INGEST, v)); - cs.addSettingsUpdateConsumer(QUERY_MAX_SETTING, v -> allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, v)); - cs.addSettingsUpdateConsumer(QUERY_MIN_SETTING, v -> allocator.setPoolMin(NativeAllocatorPoolConfig.POOL_QUERY, v)); - - // Grouped validator runs across the related settings on every dynamic update so cross-setting - // invariants (sum of pool mins ≤ root, per-pool min ≤ max) are enforced post-startup. - cs.addSettingsUpdateConsumer(s -> {}, MIN_MAX_SETTINGS, ArrowBasePlugin::validateUpdate); - } + // ─── Private helpers ───────────────────────────────────────────────────────── - private static final List> MIN_MAX_SETTINGS = List.of( - ROOT_LIMIT_SETTING, - FLIGHT_MIN_SETTING, - FLIGHT_MAX_SETTING, - INGEST_MIN_SETTING, - INGEST_MAX_SETTING, - QUERY_MIN_SETTING, - QUERY_MAX_SETTING - ); + private synchronized void startRebalancer(ArrowNativeAllocator allocator, Supplier budgetSupplier, long intervalSeconds) { + if (rebalancer != null || rebalancerScheduler != null) return; - private static void validateUpdate(Settings settings) { - long rootLimit = ROOT_LIMIT_SETTING.get(settings); - long flightMin = FLIGHT_MIN_SETTING.get(settings); - long flightMax = FLIGHT_MAX_SETTING.get(settings); - long ingestMin = INGEST_MIN_SETTING.get(settings); - long ingestMax = INGEST_MAX_SETTING.get(settings); - long queryMin = QUERY_MIN_SETTING.get(settings); - long queryMax = QUERY_MAX_SETTING.get(settings); - validateMinMax(NativeAllocatorPoolConfig.POOL_FLIGHT, flightMin, flightMax); - validateMinMax(NativeAllocatorPoolConfig.POOL_INGEST, ingestMin, ingestMax); - validateMinMax(NativeAllocatorPoolConfig.POOL_QUERY, queryMin, queryMax); - validateMinSum(rootLimit, flightMin, ingestMin, queryMin); - } + long budget = budgetSupplier.get(); + if (budget <= 0) return; + if (intervalSeconds <= 0) return; - @Override - public List> getSettings() { - return List.of( - ROOT_LIMIT_SETTING, - FLIGHT_MIN_SETTING, - FLIGHT_MAX_SETTING, - INGEST_MIN_SETTING, - INGEST_MAX_SETTING, - QUERY_MIN_SETTING, - QUERY_MAX_SETTING, - REBALANCE_INTERVAL_SETTING + NativeMemoryRebalancer nativeRebalancer = new NativeMemoryRebalancer( + allocator, + budgetSupplier, + PRESSURE_THRESHOLD_SETTING.getDefault(Settings.EMPTY), + IDLE_THRESHOLD_SETTING.getDefault(Settings.EMPTY), + SHRINK_FACTOR_SETTING.getDefault(Settings.EMPTY) ); + this.rebalancer = nativeRebalancer; + + Scheduler.SafeScheduledThreadPoolExecutor executor = new Scheduler.SafeScheduledThreadPoolExecutor(1, r -> { + Thread t = new Thread(r, "native-allocator-rebalancer"); + t.setDaemon(true); + return t; + }); + executor.setRemoveOnCancelPolicy(true); + this.rebalancerScheduler = executor; + + rebalanceTask = rebalancerScheduler.scheduleAtFixedRate(nativeRebalancer, intervalSeconds, intervalSeconds, TimeUnit.SECONDS); + } + + private synchronized void cancelRebalanceTask() { + ScheduledFuture existing = rebalanceTask; + if (existing != null) { + FutureUtils.cancel(existing); + rebalanceTask = null; + } + rebalancer = null; + if (rebalancerScheduler != null) { + rebalancerScheduler.shutdown(); + rebalancerScheduler = null; + } + } + + private void updateRebalanceInterval(long newInterval) { + cancelRebalanceTask(); + if (newInterval > 0 && rebalancerScheduler != null && rebalancer != null) { + rebalanceTask = rebalancerScheduler.scheduleAtFixedRate(rebalancer, newInterval, newInterval, TimeUnit.SECONDS); + } } private static void validateMinMax(String poolName, long min, long max) { @@ -362,36 +377,27 @@ private static void validateMinMax(String poolName, long min, long max) { } } - private static void validateMinSum(long rootLimit, long... mins) { - if (rootLimit == Long.MAX_VALUE) { - return; - } - long sum = 0; - for (long min : mins) { - try { - sum = Math.addExact(sum, min); - } catch (ArithmeticException overflow) { - throw new IllegalArgumentException("Sum of pool minimums overflows.", overflow); - } + static String derivePoolMaxDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return Long.toString(Long.MAX_VALUE); } - if (sum > rootLimit) { - throw new IllegalArgumentException( - "Sum of pool minimums (" - + sum - + " bytes) exceeds root limit (" - + rootLimit - + " bytes). " - + "Reduce pool minimums or increase " - + NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT - ); + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); + } + + static String derivePoolMinDefault(Settings settings, int percent) { + ByteSizeValue nativeLimit = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(settings); + if (nativeLimit.getBytes() <= 0) { + return "0"; } + return Long.toString(Math.max(0L, nativeLimit.getBytes() * percent / 100)); } - @Override - public void close() throws IOException { - if (allocator != null) { - allocator.close(); - allocator = null; + private static long parseNonNegativeLong(String s, String settingName) { + long v = Long.parseLong(s); + if (v < 0) { + throw new IllegalArgumentException("Setting [" + settingName + "] must be >= 0, got " + v); } + return v; } } diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java new file mode 100644 index 0000000000000..70928e93713a0 --- /dev/null +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowBaseStatsAction.java @@ -0,0 +1,75 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.plugin.stats.NativeAllocatorPoolStats; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.BytesRestResponse; +import org.opensearch.rest.RestRequest; +import org.opensearch.transport.client.node.NodeClient; + +import java.util.List; +import java.util.function.Supplier; + +/** + * REST handler exposing per-pool native memory stats at {@code _plugins/arrow_base/stats}. + */ +public class ArrowBaseStatsAction extends BaseRestHandler { + private final Supplier statsSupplier; + + /** + * Creates a new stats action. + * @param statsSupplier supplier of pool stats + */ + public ArrowBaseStatsAction(Supplier statsSupplier) { + this.statsSupplier = statsSupplier; + } + + @Override + public String getName() { + return "arrow_base_stats_action"; + } + + @Override + public List routes() { + return List.of(new Route(RestRequest.Method.GET, "_plugins/arrow_base/stats")); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + return channel -> { + NativeAllocatorPoolStats stats = statsSupplier.get(); + XContentBuilder builder = channel.newBuilder(); + builder.startObject(); + builder.startObject("memory_pools"); + if (stats != null) { + builder.startObject("runtime"); + builder.field("allocated_bytes", stats.getNativeAllocatedBytes()); + builder.field("resident_bytes", stats.getNativeResidentBytes()); + builder.endObject(); + builder.startObject("pools"); + for (NativeAllocatorPoolStats.PoolStats pool : stats.getPools()) { + builder.startObject(pool.getName()); + builder.field("allocated_bytes", pool.getAllocatedBytes()); + builder.field("peak_bytes", pool.getPeakBytes()); + builder.field("limit_bytes", pool.getLimitBytes()); + builder.field("min_bytes", pool.getMinBytes()); + builder.field("group", pool.getGroup()); + builder.endObject(); + } + builder.endObject(); + } + builder.endObject(); + builder.endObject(); + channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder)); + }; + } +} diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java index 892c43d2cb2c8..baf5ff398b760 100644 --- a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/ArrowNativeAllocator.java @@ -10,197 +10,296 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.function.Consumer; +import java.util.function.Supplier; /** * Arrow-backed implementation of {@link NativeAllocator}. * - *

Owns a single {@link RootAllocator} for the node. All plugins that need - * Arrow buffers obtain pool handles from this class via the SPI interface. - * - *

Elastic rebalancing

- *

A background task periodically redistributes unused capacity across pools. - * Each pool has a guaranteed limit (configured via settings). When other - * pools are idle, an active pool can temporarily grow beyond its guarantee up to - * the root limit. When contention rises, pools shrink back toward their guarantee. - * This prevents idle capacity from being wasted while maintaining isolation under load. - * - *

Constructed once by {@link ArrowBasePlugin#createComponents} and exposed to - * downstream plugins via Guice and {@code PluginComponentRegistry} so consumers - * receive the instance through explicit dependency injection rather than a static - * singleton. + *

Owns a single {@link RootAllocator} (set to {@code Long.MAX_VALUE} — per-pool + * limits are the real enforcement). Manages both Arrow-backed pools and virtual pools. */ public class ArrowNativeAllocator implements NativeAllocator { + private static final Logger logger = LogManager.getLogger(ArrowNativeAllocator.class); + private final RootAllocator root; private final ConcurrentMap pools = new ConcurrentHashMap<>(); - private final ConcurrentMap poolMins = new ConcurrentHashMap<>(); - private final ConcurrentMap poolMaxes = new ConcurrentHashMap<>(); - private final ScheduledExecutorService rebalancer; - private volatile ScheduledFuture rebalanceTask; - /** - * True iff the rebalancer is configured to run periodically. Used by - * {@link #getOrCreatePool} to decide each pool's initial child-allocator - * limit: when rebalancing is enabled, pools start at {@code min} and grow - * via the next rebalance tick (preserving the original PR's - * "guarantee + burst" semantics); when rebalancing is disabled, pools - * start at {@code max} so consumers can allocate immediately without - * waiting for a tick that never comes. - */ - private volatile boolean rebalancerEnabled = false; + private final ConcurrentMap virtualPools = new ConcurrentHashMap<>(); + private final ConcurrentMap poolConfigs = new ConcurrentHashMap<>(); + private final ConcurrentMap>> poolGroupLimitListeners = new ConcurrentHashMap<>(); + private final List statsRefreshers = new CopyOnWriteArrayList<>(); + private volatile Supplier nativeMemoryStatsSupplier; + private volatile long budget = Long.MAX_VALUE; /** * Creates a new allocator with a fresh RootAllocator. - * - * @param rootLimit maximum bytes for the root allocator */ - public ArrowNativeAllocator(long rootLimit) { - this.root = new RootAllocator(rootLimit); - org.opensearch.threadpool.Scheduler.SafeScheduledThreadPoolExecutor executor = - new org.opensearch.threadpool.Scheduler.SafeScheduledThreadPoolExecutor(1, r -> { - Thread t = new Thread(r, "native-allocator-rebalancer"); - t.setDaemon(true); - return t; - }); - executor.setRemoveOnCancelPolicy(true); - this.rebalancer = executor; + public ArrowNativeAllocator() { + this.root = new RootAllocator(Long.MAX_VALUE); } /** - * Schedules (or reschedules) the rebalancer at the given interval. - * A value of 0 disables rebalancing. + * Sets the total native memory budget for validation. * - * @param intervalSeconds rebalance period in seconds, or 0 to disable + * @param budget node.native_memory.limit in bytes */ - public void setRebalanceInterval(long intervalSeconds) { - ScheduledFuture existing = rebalanceTask; - if (existing != null) { - org.opensearch.common.util.concurrent.FutureUtils.cancel(existing); - rebalanceTask = null; - } - rebalancerEnabled = intervalSeconds > 0; - if (rebalancerEnabled) { - rebalanceTask = rebalancer.scheduleAtFixedRate(this::rebalance, intervalSeconds, intervalSeconds, TimeUnit.SECONDS); - } + public void setBudget(long budget) { + this.budget = budget; } - @Override - public PoolHandle getOrCreatePool(String poolName, long limit) { - return getOrCreatePool(poolName, limit, limit); - } + // ─── Public / SPI methods ─────────────────────────────────────────────────── - /** - * Creates or returns a pool with min/max limits. - * - * @param poolName logical pool name - * @param min guaranteed minimum bytes (always available) - * @param max maximum bytes the pool can burst to - * @return the pool handle - */ - public PoolHandle getOrCreatePool(String poolName, long min, long max) { - poolMins.putIfAbsent(poolName, min); - poolMaxes.putIfAbsent(poolName, max); + @Override + public PoolHandle getOrCreatePool(String poolName, long min, long max, PoolGroup group) { + validateSumMaxesWithinBudget(poolName, max); + poolConfigs.putIfAbsent(poolName, new PoolConfig(min, max, group)); return pools.computeIfAbsent(poolName, name -> { - // Pick an initial limit that's safe for both rebalancer-on and rebalancer-off - // deployments. When rebalancing is enabled, start at min (the original PR's - // "guarantee + burst" semantics): the next rebalance tick will distribute - // headroom up to each pool's max. When rebalancing is disabled (the default), - // pools with min=0 would otherwise reject every allocation until a tick that - // never comes — start at max so consumers can allocate immediately. - long initial = rebalancerEnabled ? min : max; - BufferAllocator child = root.newChildAllocator(name, 0, initial); + BufferAllocator child = root.newChildAllocator(name, 0, max); return new ArrowPoolHandle(child); }); } @Override public void setPoolLimit(String poolName, long newLimit) { - ArrowPoolHandle handle = pools.get(poolName); - if (handle == null) { - throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + PoolConfig config = poolConfigs.get(poolName); + if (config != null) { + config.max = newLimit; + } + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + arrowHandle.allocator.setLimit(newLimit); + return; + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + vp.setLimit(newLimit); + return; + } + throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + } + + @Override + public VirtualPoolHandle registerVirtualPool(String poolName, long min, long max, PoolGroup group, Consumer limitSetter) { + if (min > max) { + throw new IllegalArgumentException("Pool '" + poolName + "' min (" + min + ") exceeds max (" + max + ")"); + } + validateSumMaxesWithinBudget(poolName, max); + VirtualPoolHandleImpl handle = new VirtualPoolHandleImpl(poolName, max, limitSetter); + VirtualPoolHandleImpl existing = virtualPools.putIfAbsent(poolName, handle); + if (existing != null || pools.containsKey(poolName)) { + virtualPools.remove(poolName, handle); + throw new IllegalStateException("Pool '" + poolName + "' already registered"); } - poolMaxes.put(poolName, newLimit); - handle.allocator.setLimit(newLimit); + poolConfigs.put(poolName, new PoolConfig(min, max, group)); + limitSetter.accept(max); + return handle; + } + + @Override + public void setPoolMin(String poolName, long newMin) { + PoolConfig config = poolConfigs.get(poolName); + if (config != null) { + config.min = newMin; + } + // Raise live limit if newMin exceeds current effective limit + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + long max = config != null ? config.max : Long.MAX_VALUE; + long current = arrowHandle.allocator.getLimit(); + long target = Math.min(newMin, max); + if (target > current) { + arrowHandle.allocator.setLimit(target); + } + return; + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + long max = config != null ? config.max : Long.MAX_VALUE; + long current = vp.limit(); + long target = Math.min(newMin, max); + if (target > current) { + vp.setLimit(target); + } + } + } + + @Override + public Set getAllPoolNames() { + Set all = new HashSet<>(pools.keySet()); + all.addAll(virtualPools.keySet()); + return Collections.unmodifiableSet(all); + } + + @Override + public void addStatsRefresher(Runnable refresher) { + statsRefreshers.add(refresher); + } + + @Override + public void setNativeMemoryStatsSupplier(Supplier supplier) { + this.nativeMemoryStatsSupplier = supplier; } /** - * Updates the minimum guaranteed bytes for a pool. The new min is recorded for the - * rebalancer (which honors it as a floor on the next tick) and also pushed to the - * live {@link BufferAllocator} so the change takes effect immediately even when - * the rebalancer is disabled — the alternative was a Dynamic setting that returned - * HTTP 200 but had no observable effect. + * Sets the effective (live) limit for a pool without updating the configured max. + * Used by the rebalancer to adjust pool limits dynamically. * - *

Live propagation rules: - *

    - *
  • If {@code newMin} exceeds the pool's current limit, the limit is raised to - * {@code newMin} (capped at the configured pool max). Children of the pool - * allocator inherit the change automatically via Arrow's parent-cap check at - * allocation time, so dynamic resizes reach in-flight workloads without an - * explicit notification SPI. - *
  • If {@code newMin} is below the current limit, the limit is left alone — - * the rebalancer is the only path that shrinks live limits, so a min change - * on its own never reduces capacity in flight. - *
- * - * @param poolName the pool name - * @param newMin new minimum bytes + * @param poolName name of the pool + * @param newLimit new effective limit in bytes */ - public void setPoolMin(String poolName, long newMin) { - ArrowPoolHandle handle = pools.get(poolName); - if (handle == null) { - throw new IllegalStateException("Pool '" + poolName + "' does not exist"); + public void setPoolEffectiveLimit(String poolName, long newLimit) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + arrowHandle.allocator.setLimit(newLimit); + return; } - poolMins.put(poolName, newMin); - long max = poolMaxes.getOrDefault(poolName, Long.MAX_VALUE); - long current = handle.allocator.getLimit(); - long target = Math.min(newMin, max); - if (target > current) { - handle.allocator.setLimit(target); + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + vp.setLimit(newLimit); + return; } + throw new IllegalStateException("Pool '" + poolName + "' does not exist"); } - @Override - public void setRootLimit(long limit) { - root.setLimit(limit); + /** + * Resets all pools to their configured max. Called when the rebalancer is disabled. + * Logs a warning for any pool that was bursting above its max. + */ + public void resetAllPoolsToMax() { + for (String name : getAllPoolNames()) { + PoolConfig config = poolConfigs.get(name); + long max = config != null ? config.max : Long.MAX_VALUE; + long current = getEffectiveLimit(name); + if (current > max) { + logger.warn( + "Pool [{}] effective limit {} exceeds max {}, resetting to max. In-flight allocations may be rejected.", + name, + current, + max + ); + } + setPoolEffectiveLimit(name, max); + } } /** - * Returns a point-in-time stats snapshot across all pools. Used by the - * {@code NativeAllocatorStatsRegistry} component published from - * {@code ArrowBasePlugin.createComponents()} and wired into {@code NodeService} to - * render allocator state under {@code _nodes/stats[/native_allocator]}. + * Convenience method for plugins that have Setting objects. Registers the virtual pool + * and auto-wires dynamic setting listeners for min/max changes. + * + * @param poolName name of the virtual pool + * @param minSetting setting for minimum bytes + * @param maxSetting setting for maximum bytes + * @param settings current node settings + * @param clusterSettings cluster settings for dynamic updates + * @param group pool group assignment + * @param limitSetter callback invoked when the pool limit changes + */ + public VirtualPoolHandle registerVirtualPool( + String poolName, + Setting minSetting, + Setting maxSetting, + Settings settings, + ClusterSettings clusterSettings, + PoolGroup group, + Consumer limitSetter + ) { + long min = minSetting.get(settings); + long max = maxSetting.get(settings); + VirtualPoolHandle handle = registerVirtualPool(poolName, min, max, group, limitSetter); + + clusterSettings.addSettingsUpdateConsumer(maxSetting, newMax -> setPoolLimit(poolName, newMax)); + clusterSettings.addSettingsUpdateConsumer(minSetting, newMin -> setPoolMin(poolName, newMin)); + + return handle; + } + + /** + * Returns a point-in-time stats snapshot across all pools. */ public NativeAllocatorPoolStats stats() { + refreshStats(); + + long nativeAllocated = -1; + long nativeResident = -1; + Supplier supplier = this.nativeMemoryStatsSupplier; + if (supplier != null) { + try { + long[] stats = supplier.get(); + if (stats != null && stats.length >= 2) { + nativeAllocated = stats[0]; + nativeResident = stats[1]; + } + } catch (Exception e) { + // best-effort + } + } + List poolStats = new ArrayList<>(); for (var entry : pools.entrySet()) { BufferAllocator alloc = entry.getValue().allocator; + PoolConfig config = poolConfigs.get(entry.getKey()); poolStats.add( new NativeAllocatorPoolStats.PoolStats( entry.getKey(), alloc.getAllocatedMemory(), alloc.getPeakMemoryAllocation(), - alloc.getLimit() + alloc.getLimit(), + config != null && config.group != null ? config.group.getName() : null, + config != null ? config.min : 0L + ) + ); + } + for (var entry : virtualPools.entrySet()) { + VirtualPoolHandleImpl vp = entry.getValue(); + PoolConfig config = poolConfigs.get(entry.getKey()); + poolStats.add( + new NativeAllocatorPoolStats.PoolStats( + entry.getKey(), + vp.allocatedBytes(), + vp.peakBytes(), + vp.limit(), + config != null && config.group != null ? config.group.getName() : null, + config != null ? config.min : 0L ) ); } - return new NativeAllocatorPoolStats(root.getAllocatedMemory(), root.getPeakMemoryAllocation(), root.getLimit(), poolStats); + + return new NativeAllocatorPoolStats(nativeAllocated, nativeResident, poolStats); + } + + /** + * Runs all registered stats refreshers. + */ + public void refreshStats() { + for (Runnable refresher : statsRefreshers) { + try { + refresher.run(); + } catch (Exception e) { + // best-effort + } + } } @Override public void close() { - rebalancer.shutdownNow(); pools.forEach((name, handle) -> { try { handle.allocator.close(); @@ -209,72 +308,23 @@ public void close() { } }); pools.clear(); - // Close any remaining child allocators (e.g., ad-hoc children created via ArrowAllocatorService) + virtualPools.clear(); for (BufferAllocator child : new ArrayList<>(root.getChildAllocators())) { try { child.close(); } catch (Exception e) { - // best-effort — log but don't block shutdown + // best-effort } } root.close(); } - /** - * Redistributes unused capacity across pools based on min/max guarantees. - * - *

Algorithm: - *

    - *
  1. Every pool is guaranteed at least its configured min
    2. - *
  2. Compute headroom = rootLimit - sum(all pool current allocations)
    3. - *
  3. Distribute headroom equally across all pools (not just active ones), capped - * at each pool's max. Distributing to all pools — including those with zero - * current allocation — avoids the dead-pool corner case where a pool with - * min = 0 starts at limit = 0, can never make its first allocation, and so - * never becomes "active" enough to receive a bonus. Pools that don't need the - * headroom stay at min naturally because their max caps the bonus.
    4. - *
  4. No pool's limit ever drops below its current allocation or its min
    5. - *
- */ - void rebalance() { - if (pools.isEmpty()) return; - - long rootLimit = root.getLimit(); - long totalAllocated = 0; - - for (Map.Entry entry : pools.entrySet()) { - totalAllocated += entry.getValue().allocator.getAllocatedMemory(); - } - - long headroom = Math.max(0, rootLimit - totalAllocated); - int poolCount = pools.size(); - long bonusPerPool = poolCount > 0 ? headroom / poolCount : 0; - - for (Map.Entry entry : pools.entrySet()) { - String name = entry.getKey(); - BufferAllocator alloc = entry.getValue().allocator; - long min = poolMins.getOrDefault(name, 0L); - long max = poolMaxes.getOrDefault(name, Long.MAX_VALUE); - long currentAllocation = alloc.getAllocatedMemory(); - - long effectiveLimit = min + bonusPerPool; - - // Cap at pool's max - effectiveLimit = Math.min(effectiveLimit, max); - // Never drop below current allocation or min - effectiveLimit = Math.max(effectiveLimit, currentAllocation); - effectiveLimit = Math.max(effectiveLimit, min); - // Never exceed root - effectiveLimit = Math.min(effectiveLimit, rootLimit); - - alloc.setLimit(effectiveLimit); - } - } + // ─── Package-private accessors (used by rebalancer and tests) ──────────────── /** * Returns the underlying Arrow allocator for a pool. * - * @param poolName name of the pool to look up + * @param poolName name of the pool */ public BufferAllocator getPoolAllocator(String poolName) { ArrowPoolHandle handle = pools.get(poolName); @@ -284,16 +334,12 @@ public BufferAllocator getPoolAllocator(String poolName) { return handle.allocator; } - /** - * Returns the root Arrow allocator. - */ + /** Returns the root Arrow allocator. */ public BufferAllocator getRootAllocator() { return root; } - /** - * Returns all registered pool names. - */ + /** Returns all registered pool names (Arrow pools only). */ public Set getPoolNames() { return Collections.unmodifiableSet(pools.keySet()); } @@ -304,7 +350,8 @@ public Set getPoolNames() { * @param poolName name of the pool */ public long getPoolMin(String poolName) { - return poolMins.getOrDefault(poolName, 0L); + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.min : 0L; } /** @@ -313,7 +360,206 @@ public long getPoolMin(String poolName) { * @param poolName name of the pool */ public long getPoolMax(String poolName) { - return poolMaxes.getOrDefault(poolName, Long.MAX_VALUE); + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.max : Long.MAX_VALUE; + } + + /** + * Returns the group for a pool, or null if not assigned. + * + * @param poolName name of the pool + */ + public PoolGroup getPoolGroup(String poolName) { + PoolConfig config = poolConfigs.get(poolName); + return config != null ? config.group : null; + } + + /** + * Returns the allocated bytes for a virtual pool. + * + * @param poolName name of the virtual pool + */ + public long getVirtualPoolAllocated(String poolName) { + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + return vp != null ? vp.allocatedBytes() : 0; + } + + /** + * Returns the current limit for a virtual pool. + * + * @param poolName name of the virtual pool + */ + public long getVirtualPoolLimit(String poolName) { + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + return vp != null ? vp.limit() : 0; + } + + /** + * Returns the effective limit for any pool (Arrow or virtual). + * + * @param poolName name of the pool + */ + public long getEffectiveLimit(String poolName) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + return arrowHandle.allocator.getLimit(); + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + return vp.limit(); + } + return 0; + } + + /** + * Returns the allocated bytes for any pool (Arrow or virtual). + * + * @param poolName name of the pool + */ + public long getAllocated(String poolName) { + ArrowPoolHandle arrowHandle = pools.get(poolName); + if (arrowHandle != null) { + return arrowHandle.allocator.getAllocatedMemory(); + } + VirtualPoolHandleImpl vp = virtualPools.get(poolName); + if (vp != null) { + return vp.allocatedBytes(); + } + return 0; + } + + /** Returns the native memory stats supplier. */ + public Supplier getNativeMemoryStatsSupplier() { + return nativeMemoryStatsSupplier; + } + + /** + * Registers a listener invoked when the grouped effective limit for the given pool group changes. + * The consumer receives the new total effective limit (sum of all pools in the group). + * + * @param group the pool group to listen on + * @param listener consumer that receives the new grouped limit in bytes + */ + public void addPoolGroupLimitListener(PoolGroup group, Consumer listener) { + poolGroupLimitListeners.computeIfAbsent(group, k -> new CopyOnWriteArrayList<>()).add(listener); + } + + /** + * Fires listeners for the specified pool group with the current grouped effective limit sum. + * Called by the rebalancer after all pool limits for a tick have been updated. + * + * @param group the pool group whose listeners should be notified + */ + public void firePoolGroupListeners(PoolGroup group) { + List> listeners = poolGroupLimitListeners.get(group); + if (listeners == null || listeners.isEmpty()) return; + + long groupSum = 0; + for (var entry : poolConfigs.entrySet()) { + if (entry.getValue().group == group) { + groupSum += getEffectiveLimit(entry.getKey()); + } + } + long finalSum = groupSum; + for (Consumer listener : listeners) { + try { + listener.accept(finalSum); + } catch (Exception e) { + logger.warn( + () -> new org.apache.logging.log4j.message.ParameterizedMessage( + "Pool group limit listener failed for group [{}]", + group + ), + e + ); + } + } + } + + // ─── Private helpers ───────────────────────────────────────────────────────── + + private void validateSumMaxesWithinBudget(String newPoolName, long newPoolMax) { + if (budget == Long.MAX_VALUE || budget <= 0) { + return; + } + long sumMaxes = newPoolMax; + for (var entry : poolConfigs.entrySet()) { + if (entry.getKey().equals(newPoolName) == false) { + sumMaxes += entry.getValue().max; + } + } + if (sumMaxes > budget) { + throw new IllegalArgumentException( + "Sum of pool max limits (" + + sumMaxes + + " bytes) exceeds native memory budget (" + + budget + + " bytes). Reduce pool max settings or increase the budget." + ); + } + } + + // ─── Inner classes ─────────────────────────────────────────────────────────── + + /** + * Mutable configuration for a pool: min, max, and group. + */ + static class PoolConfig { + volatile long min; + volatile long max; + final PoolGroup group; + + PoolConfig(long min, long max, PoolGroup group) { + this.min = min; + this.max = max; + this.group = group; + } + } + + /** + * Virtual pool handle implementation. Tracks stats reported from native layer + * and delegates limit changes to the registered callback. + */ + public static class VirtualPoolHandleImpl implements VirtualPoolHandle { + private final String name; + private volatile long limit; + private volatile long allocatedBytes; + private volatile long peakBytes; + private final Consumer limitSetter; + + VirtualPoolHandleImpl(String name, long limit, Consumer limitSetter) { + this.name = name; + this.limit = limit; + this.limitSetter = limitSetter; + } + + @Override + public void updateStats(long allocated, long peak) { + this.allocatedBytes = allocated; + this.peakBytes = peak; + } + + void setLimit(long newLimit) { + this.limit = newLimit; + if (limitSetter != null) { + limitSetter.accept(newLimit); + } + } + + @Override + public long allocatedBytes() { + return allocatedBytes; + } + + @Override + public long peakBytes() { + return peakBytes; + } + + @Override + public long limit() { + return limit; + } } /** diff --git a/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java new file mode 100644 index 0000000000000..5a5294b2cb4b0 --- /dev/null +++ b/plugins/arrow-base/src/main/java/org/opensearch/arrow/allocator/NativeMemoryRebalancer.java @@ -0,0 +1,241 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.opensearch.arrow.spi.PoolGroup; + +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; + +/** + * Periodic rebalancer that redistributes native memory across pools. + * + *

Algorithm: + *

    + *
  • Pools start at their configured max on registration
    • + *
  • If no pool is under pressure, return early (no-op)
    • + *
  • Idle pools (utilization < idle_threshold) are shrunk, never below min
    • + *
  • Pressured pools (utilization > pressure_threshold) receive freed capacity, can exceed max
    • + *
  • Excess freed capacity is returned to idle pools proportionally
    • + *
  • Invariant: sum(effective_limits) <= budget at all times
    • + *
+ * + * @opensearch.internal + */ +public class NativeMemoryRebalancer implements Runnable { + + private static final Logger logger = LogManager.getLogger(NativeMemoryRebalancer.class); + + private final ArrowNativeAllocator allocator; + private final Supplier budgetSupplier; + + private volatile double pressureThreshold; + private volatile double idleThreshold; + private volatile double shrinkFactor; + + /** + * Creates a new rebalancer. + * + * @param allocator the allocator managing all pools + * @param budgetSupplier supplies the current budget value + * @param pressureThreshold utilization above this triggers growth (default 0.75) + * @param idleThreshold utilization below this means pool can give back capacity (default 0.50) + * @param shrinkFactor factor to shrink idle pools by — new limit = limit * (1 - shrinkFactor) (default 0.10) + */ + public NativeMemoryRebalancer( + ArrowNativeAllocator allocator, + Supplier budgetSupplier, + double pressureThreshold, + double idleThreshold, + double shrinkFactor + ) { + this.allocator = allocator; + this.budgetSupplier = budgetSupplier; + this.pressureThreshold = pressureThreshold; + this.idleThreshold = idleThreshold; + this.shrinkFactor = shrinkFactor; + } + + /** + * Updates the pressure threshold dynamically. + * + * @param value new threshold (0.0 to 1.0) + */ + public void setPressureThreshold(double value) { + this.pressureThreshold = value; + } + + /** + * Updates the idle threshold dynamically. + * + * @param value new threshold (0.0 to 1.0) + */ + public void setIdleThreshold(double value) { + this.idleThreshold = value; + } + + /** + * Updates the shrink factor dynamically. + * + * @param value new factor (0.0 to 1.0) + */ + public void setShrinkFactor(double value) { + this.shrinkFactor = value; + } + + @Override + public void run() { + try { + rebalance(); + } catch (Exception e) { + logger.warn("Rebalancer tick failed", e); + } + } + + void rebalance() { + Set allPools = allocator.getAllPoolNames(); + if (allPools.isEmpty()) return; + + long budget = budgetSupplier.get(); + if (budget <= 0 || budget == Long.MAX_VALUE) return; + + // Refresh stats from native layers + allocator.refreshStats(); + + // Snapshot per-pool state + Map snapshots = new HashMap<>(); + for (String name : allPools) { + long allocated = allocator.getAllocated(name); + long effectiveLimit = allocator.getEffectiveLimit(name); + long min = allocator.getPoolMin(name); + long max = allocator.getPoolMax(name); + double utilization = effectiveLimit > 0 ? (double) allocated / effectiveLimit : 0; + snapshots.put(name, new PoolSnapshot(allocated, effectiveLimit, min, max, utilization)); + } + + // Identify pressured pools — if none, nothing to do + Map desires = new HashMap<>(); + long totalDesired = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization > pressureThreshold) { + long desired = Math.max(1, (long) (s.allocated * 0.25)); + desires.put(entry.getKey(), desired); + totalDesired += desired; + } + } + if (totalDesired == 0) { + logger.debug("Rebalancer: no pools under pressure, skipping"); + return; + } + + // Shrink idle pools, floor at min + long freedCapacity = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization < idleThreshold) { + long newLimit = Math.max((long) (s.effectiveLimit * (1.0 - shrinkFactor)), s.min); + newLimit = Math.max(newLimit, s.allocated); + if (newLimit < s.effectiveLimit) { + freedCapacity += s.effectiveLimit - newLimit; + allocator.setPoolEffectiveLimit(entry.getKey(), newLimit); + s.effectiveLimit = newLimit; + } + } + } + + if (freedCapacity == 0) { + logger.debug("Rebalancer: no capacity freed from idle pools"); + return; + } + + // Distribute freed capacity to pressured pools (can exceed max) + long totalGranted = 0; + long grantCap = Math.min(freedCapacity, totalDesired); + for (var entry : desires.entrySet()) { + String name = entry.getKey(); + long desired = entry.getValue(); + PoolSnapshot s = snapshots.get(name); + long grant = (long) ((double) grantCap * desired / totalDesired); + grant = Math.min(grant, grantCap - totalGranted); + if (grant > 0) { + try { + long newLimit = s.effectiveLimit + grant; + allocator.setPoolEffectiveLimit(name, newLimit); + totalGranted += grant; + logger.debug("Rebalancer: grew pool [{}] by {} bytes to {} (max={})", name, grant, newLimit, s.max); + } catch (Exception e) { + logger.warn(() -> new ParameterizedMessage("Rebalancer: failed to grow pool [{}]", name), e); + } + } + } + + // Return any excess freed capacity back to idle pools + long excess = freedCapacity - totalGranted; + if (excess > 0) { + returnToIdlePools(snapshots, excess); + } + + // Notify pool group listeners after all limit changes are complete + for (PoolGroup group : PoolGroup.values()) { + allocator.firePoolGroupListeners(group); + } + } + + // ─── Private helpers ───────────────────────────────────────────────────────── + + private void returnToIdlePools(Map snapshots, long capacity) { + long totalIdleSize = 0; + for (PoolSnapshot s : snapshots.values()) { + if (s.utilization < idleThreshold) { + totalIdleSize += s.effectiveLimit; + } + } + if (totalIdleSize == 0) return; + + long totalReturned = 0; + for (var entry : snapshots.entrySet()) { + PoolSnapshot s = entry.getValue(); + if (s.utilization < idleThreshold) { + long share = (long) ((double) capacity * s.effectiveLimit / totalIdleSize); + share = Math.min(share, capacity - totalReturned); + if (share > 0) { + long newLimit = s.effectiveLimit + share; + allocator.setPoolEffectiveLimit(entry.getKey(), newLimit); + s.effectiveLimit = newLimit; + totalReturned += share; + } + } + } + } + + /** + * Point-in-time snapshot of a pool's state during one rebalance tick. + */ + static class PoolSnapshot { + final long allocated; + long effectiveLimit; + final long min; + final long max; + final double utilization; + + PoolSnapshot(long allocated, long effectiveLimit, long min, long max, double utilization) { + this.allocated = allocated; + this.effectiveLimit = effectiveLimit; + this.min = min; + this.max = max; + this.utilization = utilization; + } + } +} diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java index ad72b70b8cbbb..df2673a5a8a7b 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowBasePluginTests.java @@ -8,12 +8,11 @@ package org.opensearch.arrow.allocator; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.OutOfMemoryException; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.node.resource.tracker.ResourceTrackerSettings; import org.opensearch.test.OpenSearchTestCase; import java.util.HashSet; @@ -21,90 +20,23 @@ public class ArrowBasePluginTests extends OpenSearchTestCase { - public void testDeriveRootLimitDefaultUnsetReturnsLongMaxValue() { - // Explicit 0 expresses "AC unconfigured" — default is now ram - heap, so Settings.EMPTY - // would resolve to a real value on whatever machine the test runs on. - Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); - assertEquals(Long.toString(Long.MAX_VALUE), ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testDeriveRootLimitDefaultUsesAcLimitWhenSet() { - Settings s = Settings.builder().put("node.native_memory.limit", "1gb").build(); - // ROOT_LIMIT defaults to 20% of node.native_memory.limit — the Arrow framework gets a - // small fraction of native budget; DataFusion's Rust runtime takes the larger share. - long oneGiB = 1024L * 1024 * 1024; - assertEquals(Long.toString(oneGiB * 20 / 100), ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testDeriveRootLimitDefaultIgnoresBufferPercent() { - // node.native_memory.buffer_percent is admission control's throttle margin, not a - // framework budget reduction. The framework default takes its 20% fraction off - // node.native_memory.limit directly so AC's safety margin sits between AC's throttle - // threshold and the framework's hard cap rather than being collapsed into the cap. - // 1000 bytes limit, 20% buffer => root.limit still 20% of 1000 = 200. - Settings s = Settings.builder().put("node.native_memory.limit", "1000b").put("node.native_memory.buffer_percent", 20).build(); - assertEquals("200", ArrowBasePlugin.deriveRootLimitDefault(s)); - } - - public void testRootLimitSettingExposesDerivedDefault() { - Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); - // 20% of 10 GiB. - long expected = 10L * 1024 * 1024 * 1024 * 20 / 100; - assertEquals(Long.valueOf(expected), ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - } - - public void testRootLimitSettingExplicitOverridesDerived() { - Settings s = Settings.builder() - .put("node.native_memory.limit", "8gb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 1024L) - .build(); - assertEquals(Long.valueOf(1024L), ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - } - - public void testRootLimitRejectsNegative() { - Settings s = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, -1L).build(); - IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> ArrowBasePlugin.ROOT_LIMIT_SETTING.get(s)); - assertTrue(e.getMessage().contains("must be >= 0")); - } - public void testQuerySettingsExposeDefaults() { // Explicit 0 expresses "AC unconfigured" so QUERY_MAX falls back to Long.MAX_VALUE. - // Settings.EMPTY would resolve via ram - heap default to a finite, machine-dependent value. Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); assertEquals(Long.valueOf(0L), ArrowBasePlugin.QUERY_MIN_SETTING.get(s)); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testFlightAndIngestMinDefaultsToZero() { - // The grouped validator (validateMinSum) treats per-pool mins as a guarantee - // floor — defaults of Long.MAX_VALUE caused the validator to reject any PUT - // that set a non-MAX root. Pool mins must default to zero so the baseline - // configuration is consistent. - Settings s = Settings.EMPTY; - assertEquals(Long.valueOf(0L), ArrowBasePlugin.FLIGHT_MIN_SETTING.get(s)); - assertEquals(Long.valueOf(0L), ArrowBasePlugin.INGEST_MIN_SETTING.get(s)); - } - - public void testQuerySettingsAcceptValues() { - Settings s = Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 100L) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1000L) - .build(); - assertEquals(Long.valueOf(100L), ArrowBasePlugin.QUERY_MIN_SETTING.get(s)); - assertEquals(Long.valueOf(1000L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); + public void testFlightAndIngestMinDerivedFromBudget() { + // With node.native_memory.limit set, mins derive as percentages + Settings s = Settings.builder().put("node.native_memory.limit", "1gb").build(); + long budget = 1024L * 1024 * 1024; + // flight min = 2% of budget, ingest min = 4% of budget + assertEquals(Long.valueOf(budget * 2 / 100), ArrowBasePlugin.FLIGHT_MIN_SETTING.get(s)); + assertEquals(Long.valueOf(budget * 4 / 100), ArrowBasePlugin.INGEST_MIN_SETTING.get(s)); } - // -- Pool max defaults derived from node.native_memory.limit ---------- - // Pool maxes anchor to the operator's off-heap budget (node.native_memory.limit), - // not to native.allocator.root.limit. This matches the PR #21732 partitioning - // diagram where pool fractions (5%/8%/5%) are of native_memory.limit. Sum of - // pool maxes (18% of native_memory.limit) fits within root.limit (20% of - // native_memory.limit) by default, leaving 2 pp headroom inside the root cap. - public void testPoolMaxDefaultsAreLongMaxValueWhenAcUnset() { - // AC explicitly unconfigured — pool maxes default to Long.MAX_VALUE (unbounded), - // preserving pre-AC behaviour. The default for node.native_memory.limit is - // 79% of (ram - heap), so to test the "unset" branch we must explicitly set it to 0. Settings s = Settings.builder().put("node.native_memory.limit", "0b").build(); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); assertEquals(Long.valueOf(Long.MAX_VALUE), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); @@ -112,10 +44,6 @@ public void testPoolMaxDefaultsAreLongMaxValueWhenAcUnset() { } public void testPoolMaxDefaultsScaleFromAcBudget() { - // 10 GiB native memory limit. Pool maxes per the partitioning model in PR #21732: - // FLIGHT_MAX = 5% INGEST_MAX = 8% QUERY_MAX = 5% - // Anchored to node.native_memory.limit, not to root.limit (which defaults to 20% - // of native_memory.limit) — see derivePoolMaxDefault Javadoc. Settings s = Settings.builder().put("node.native_memory.limit", "10gb").build(); long limit = 10L * 1024 * 1024 * 1024; assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); @@ -123,51 +51,14 @@ public void testPoolMaxDefaultsScaleFromAcBudget() { assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testPoolMaxDefaultsIgnoreRootLimitOverride() { - // Pool maxes anchor to node.native_memory.limit, not to root.limit. An operator - // who overrides root.limit (e.g. to 4 GiB instead of the default 20% of - // native_memory.limit = 2 GiB) does not shrink pool defaults proportionally; - // the diagrammed partitioning of native_memory.limit holds. - Settings s = Settings.builder() - .put("node.native_memory.limit", "10gb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 4L * 1024 * 1024 * 1024) - .build(); - long limit = 10L * 1024 * 1024 * 1024; - assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(limit * 8 / 100), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(limit * 5 / 100), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); - } - public void testPoolMaxDefaultsIgnoreBufferPercent() { - // node.native_memory.buffer_percent is AC's throttle margin, not a framework budget - // reduction. Pool maxes default off node.native_memory.limit directly so AC's safety - // margin sits between AC's throttle threshold and the framework's hard cap rather than - // being collapsed into the cap. - // 1000 bytes limit, 20% buffer => pool maxes are still 5/8/5% of 1000 = 50/80/50. Settings s = Settings.builder().put("node.native_memory.limit", "1000b").put("node.native_memory.buffer_percent", 20).build(); assertEquals(Long.valueOf(50L), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); assertEquals(Long.valueOf(80L), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); assertEquals(Long.valueOf(50L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); } - public void testPoolMaxExplicitOverridesDerived() { - // Operator-set values must win over derived defaults. - Settings s = Settings.builder() - .put("node.native_memory.limit", "10gb") - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 7L) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 8L) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 9L) - .build(); - assertEquals(Long.valueOf(7L), ArrowBasePlugin.FLIGHT_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(8L), ArrowBasePlugin.INGEST_MAX_SETTING.get(s)); - assertEquals(Long.valueOf(9L), ArrowBasePlugin.QUERY_MAX_SETTING.get(s)); - } - public void testPoolMaxRejectsNegative() { - // Negative pool max is rejected at parse time, mirroring ROOT_LIMIT_SETTING. - // Each pool's parser has its own message so we exercise all three to lock down - // the per-pool error contract (and keep coverage honest on what is otherwise - // boilerplate-but-distinct branches). Settings flight = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, -1L).build(); IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> ArrowBasePlugin.FLIGHT_MAX_SETTING.get(flight)); assertTrue(e.getMessage().contains("must be >= 0")); @@ -188,253 +79,75 @@ public void testPoolMaxRejectsNegative() { } // ----------------------------------------------------------------- - // End-to-end wiring tests — verify that Setting.Property.Dynamic settings - // actually flow through to the live allocator. These guard against the - // "dynamic in name only" failure mode where a setting parses, the validator - // runs, the cluster-state update succeeds, and the runtime component - // silently does nothing because the addSettingsUpdateConsumer line was - // never registered. Bare-setter unit tests do not catch this; tests must - // drive a real ClusterSettings#applySettings round-trip. + // End-to-end wiring tests // ----------------------------------------------------------------- - /** - * Builds a {@link ClusterSettings} preloaded with all of {@link ArrowBasePlugin}'s - * settings, mirroring what {@code SettingsModule} does at node startup. Returns the - * fresh allocator with the framework's pools created and consumers registered - * — the same wiring path {@code createComponents} runs. - */ - private static ArrowNativeAllocator newWiredAllocator(Settings nodeSettings, ClusterSettings cs) { - long rootLimit = ArrowBasePlugin.ROOT_LIMIT_SETTING.get(nodeSettings); - ArrowNativeAllocator allocator = new ArrowNativeAllocator(rootLimit); - allocator.setRebalanceInterval(ArrowBasePlugin.REBALANCE_INTERVAL_SETTING.get(nodeSettings)); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_FLIGHT, - ArrowBasePlugin.FLIGHT_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.FLIGHT_MAX_SETTING.get(nodeSettings) - ); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_INGEST, - ArrowBasePlugin.INGEST_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.INGEST_MAX_SETTING.get(nodeSettings) - ); - allocator.getOrCreatePool( - NativeAllocatorPoolConfig.POOL_QUERY, - ArrowBasePlugin.QUERY_MIN_SETTING.get(nodeSettings), - ArrowBasePlugin.QUERY_MAX_SETTING.get(nodeSettings) - ); - ArrowBasePlugin.registerSettingsUpdateConsumers(cs, allocator); - return allocator; - } - private static ClusterSettings newClusterSettings(Settings nodeSettings) { Set> registered = new HashSet<>(); registered.addAll(new ArrowBasePlugin().getSettings()); return new ClusterSettings(nodeSettings, registered); } - public void testBuildAllocatorWiresAllPoolsAndSettingsConsumers() { - // Verifies the full createComponents code path — the helper extracted from - // createComponents builds the allocator, creates all three pools (FLIGHT, INGEST, - // QUERY), and registers the cluster-settings update consumers. We bypass the - // heavyweight ClusterService fixture and inject a real ClusterSettings directly, - // which is what production wiring also passes through to buildAllocator after - // unpacking the createComponents arguments. + public void testBuildAllocatorWiresAllPools() throws Exception { Settings nodeSettings = Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) + .put("node.native_memory.limit", "10gb") .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 1L * 1024 * 1024 * 1024) .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 2L * 1024 * 1024 * 1024) .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1L * 1024 * 1024 * 1024) + .put("native.allocator.rebalancer.enabled", false) .build(); ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = ArrowBasePlugin.buildAllocator(nodeSettings, cs); + ArrowBasePlugin plugin = new ArrowBasePlugin(); + long budget = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(nodeSettings).getBytes(); + ArrowNativeAllocator allocator = plugin.buildAllocator(nodeSettings, cs, () -> budget); try { - // All three pools created. Set poolNames = allocator.getPoolNames(); assertEquals("buildAllocator must register exactly the framework's three pools", 3, poolNames.size()); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_FLIGHT)); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_INGEST)); assertTrue(poolNames.contains(NativeAllocatorPoolConfig.POOL_QUERY)); - // Pool maxes match the operator-set values (rebalancer disabled by default, + // Pool maxes match the operator-set values (rebalancer disabled, // so initial limit == max). assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit()); assertEquals(2L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit()); assertEquals(1L * 1024 * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit()); - - // Cluster-settings update consumers are registered: a PUT to a pool max must - // propagate to the live allocator. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 4L * 1024 * 1024 * 1024) - .build() - ); - assertEquals( - "buildAllocator must wire the INGEST_MAX cluster-settings consumer", - 4L * 1024 * 1024 * 1024, - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit() - ); - } finally { - allocator.close(); - } - } - - public void testQueryMaxClusterSettingPropagatesToAllocator() { - // The full wired path: node starts at default settings, plugin registers - // consumers, operator PUTs a new max via _cluster/settings. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1024L * 1024 * 1024) - .build() - ); - assertEquals( - "PUT to query max must update the live BufferAllocator limit", - 1024L * 1024 * 1024, - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit() - ); - assertEquals(1024L * 1024 * 1024, allocator.getPoolMax(NativeAllocatorPoolConfig.POOL_QUERY)); - } finally { - allocator.close(); - } - } - - public void testFlightMinClusterSettingPropagatesToAllocator() { - // Min is the regression-prone path: prior to the live-propagation fix, - // setPoolMin only updated the poolMins map and operators got HTTP 200 with - // no observable behavior change. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - // Pool starts at max (rebalancer disabled by default), so a min PUT below - // the current limit is a no-op on the live limit but updates poolMins. - // Use a min ABOVE the current limit to force the live raise path. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 4L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 2L * 1024 * 1024 * 1024) - .build() - ); - assertEquals( - "PUT to flight min must update the recorded min for the rebalancer", - 2L * 1024 * 1024 * 1024, - allocator.getPoolMin(NativeAllocatorPoolConfig.POOL_FLIGHT) - ); - assertTrue( - "PUT to flight min must raise the live BufferAllocator limit when min exceeds current", - allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit() >= 2L * 1024 * 1024 * 1024 - ); - } finally { - allocator.close(); - } - } - - public void testRootLimitClusterSettingPropagatesToAllocator() { - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - cs.applySettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 16L * 1024 * 1024 * 1024).build()); - assertEquals( - "PUT to root limit must update the RootAllocator's limit", - 16L * 1024 * 1024 * 1024, - allocator.getRootAllocator().getLimit() - ); } finally { allocator.close(); + plugin.close(); } } - public void testValidatorRejectsSumOfMinsExceedingRoot() { - // The cross-setting grouped validator must reject PUTs that would over- - // subscribe the root. Test the rejection path end-to-end through ClusterSettings. - // Set node.native_memory.limit=0b explicitly so pool maxes default to Long.MAX_VALUE - // — minmax path firing first. + public void testBuildAllocatorWithRebalancerPoolsStartAtMax() throws Exception { Settings nodeSettings = Settings.builder() - .put("node.native_memory.limit", "0b") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 10L * 1024 * 1024 * 1024) + .put("node.native_memory.limit", "1gb") + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 10L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 200L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 20L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 200L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 10L * 1024 * 1024) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 200L * 1024 * 1024) + .put("native.allocator.rebalancer.enabled", true) + .put("native.allocator.rebalance.interval_seconds", 5L) .build(); ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); + + ArrowBasePlugin plugin = new ArrowBasePlugin(); + long budget2 = ResourceTrackerSettings.NODE_NATIVE_MEMORY_LIMIT_SETTING.get(nodeSettings).getBytes(); + ArrowNativeAllocator allocator = plugin.buildAllocator(nodeSettings, cs, () -> budget2); try { - // root=10gb, flight_min=6gb, ingest_min=6gb => sum_mins=12gb > root=10gb. - IllegalArgumentException e = expectThrows( - IllegalArgumentException.class, - () -> cs.applySettings( - Settings.builder() - .put("node.native_memory.limit", "0b") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 10L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 6L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 6L * 1024 * 1024 * 1024) - .build() - ) - ); - assertTrue( - "expected sum-exceeds-root in error, got: " + e.getMessage(), - e.getMessage().contains("exceeds root limit") || e.getMessage().contains("Sum of pool minimums") - ); + // Pools always start at max regardless of rebalancer state + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT).getLimit()); + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST).getLimit()); + assertEquals(200L * 1024 * 1024, allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY).getLimit()); } finally { allocator.close(); + plugin.close(); } } - public void testChildAllocatorInheritsParentCapAfterPoolLimitUpdate() { - // Sanity check for the AnalyticsSearchService / FlightTransport pattern: - // when a consumer creates a child of the framework's pool with Long.MAX_VALUE - // limit, a PUT to the pool's max takes effect on the child's allocations - // automatically via Arrow's parent-cap check at allocateBytes — no listener needed. - // - // The contract we rely on (Arrow Accountant.allocate, lines 191-203 in 18.3.0): - // when the child's reservation is exhausted, it calls parent.allocate(...) which - // checks the parent's allocationLimit on every allocation. Setting the child's own - // limit to Long.MAX_VALUE means the child has no own-cap on top of the parent's; - // setLimit on the parent is observed atomically by all subsequent allocations - // through any descendant. - Settings nodeSettings = Settings.builder().put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024).build(); - ClusterSettings cs = newClusterSettings(nodeSettings); - ArrowNativeAllocator allocator = newWiredAllocator(nodeSettings, cs); - try { - BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - BufferAllocator child = queryPool.newChildAllocator("consumer", 0, Long.MAX_VALUE); - try { - // Step 1: a small allocation through the child succeeds with the original pool max. - try (var buf = child.buffer(1024)) { - assertEquals("child accounting reflects allocation", 1024L, child.getAllocatedMemory()); - assertEquals("parent pool sees child allocation", 1024L, queryPool.getAllocatedMemory()); - } - - // Step 2: PUT a small pool max via cluster settings. - cs.applySettings( - Settings.builder() - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 8L * 1024 * 1024 * 1024) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 1L * 1024 * 1024) // 1 MB - .build() - ); - assertEquals("pool's own limit reflects the PUT", 1L * 1024 * 1024, queryPool.getLimit()); - assertEquals("child's own limit is intentionally uncapped", Long.MAX_VALUE, child.getLimit()); - - // Step 3: allocations within the new parent cap still work. - try (var withinCap = child.buffer(512 * 1024)) { // 512 KB, under 1 MB cap - assertEquals(512L * 1024, child.getAllocatedMemory()); - } - - // Step 4: allocation exceeding the new parent cap fails — this is the - // behavior the deleted listener pattern was emulating, now provided - // natively by Arrow's parent-cap check. - expectThrows(OutOfMemoryException.class, () -> child.buffer(2L * 1024 * 1024)); // 2 MB, over 1 MB cap - } finally { - child.close(); - } - } finally { - allocator.close(); - } + public void testRebalanceIntervalSettingDefault() { + assertEquals(Long.valueOf(5L), ArrowBasePlugin.REBALANCE_INTERVAL_SETTING.get(Settings.EMPTY)); } } diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java index dc10a93bd6b74..c9e370bc22413 100644 --- a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/ArrowNativeAllocatorTests.java @@ -10,6 +10,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.opensearch.arrow.spi.NativeAllocator; +import org.opensearch.arrow.spi.PoolGroup; import org.opensearch.plugin.stats.NativeAllocatorPoolStats; import org.opensearch.test.OpenSearchTestCase; @@ -20,7 +21,7 @@ public class ArrowNativeAllocatorTests extends OpenSearchTestCase { @Override public void setUp() throws Exception { super.setUp(); - allocator = new ArrowNativeAllocator(1024L * 1024 * 1024); // 1 GB for tests + allocator = new ArrowNativeAllocator(); } @Override @@ -30,21 +31,21 @@ public void tearDown() throws Exception { } public void testCreatePool() { - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("test-pool", 100 * 1024 * 1024); + NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("test-pool", 0L, 100 * 1024 * 1024, null); assertNotNull(handle); assertEquals(100 * 1024 * 1024, handle.limit()); assertEquals(0, handle.allocatedBytes()); } public void testGetOrCreatePoolIdempotent() { - NativeAllocator.PoolHandle first = allocator.getOrCreatePool("idempotent", 50 * 1024 * 1024); - NativeAllocator.PoolHandle second = allocator.getOrCreatePool("idempotent", 999 * 1024 * 1024); + NativeAllocator.PoolHandle first = allocator.getOrCreatePool("idempotent", 0L, 50 * 1024 * 1024, null); + NativeAllocator.PoolHandle second = allocator.getOrCreatePool("idempotent", 0L, 999 * 1024 * 1024, null); assertSame(first, second); assertEquals(50 * 1024 * 1024, second.limit()); } public void testPoolChildAllocation() { - allocator.getOrCreatePool("parent", 200 * 1024 * 1024); + allocator.getOrCreatePool("parent", 0L, 200 * 1024 * 1024, null); BufferAllocator child = allocator.getPoolAllocator("parent").newChildAllocator("child-1", 0, 50 * 1024 * 1024); try { child.buffer(1024).close(); @@ -55,7 +56,7 @@ public void testPoolChildAllocation() { } public void testSetPoolLimit() { - allocator.getOrCreatePool("resizable", 100 * 1024 * 1024); + allocator.getOrCreatePool("resizable", 0L, 100 * 1024 * 1024, null); allocator.setPoolLimit("resizable", 200 * 1024 * 1024); assertEquals(200 * 1024 * 1024, allocator.getPoolAllocator("resizable").getLimit()); } @@ -68,147 +69,54 @@ public void testGetPoolAllocatorNonExistent() { expectThrows(IllegalStateException.class, () -> allocator.getPoolAllocator("ghost")); } - public void testSetRootLimit() { - allocator.setRootLimit(512 * 1024 * 1024); - assertEquals(512 * 1024 * 1024, allocator.getRootAllocator().getLimit()); - } - public void testStats() { - allocator.getOrCreatePool("stats-pool", 64 * 1024 * 1024); + allocator.getOrCreatePool("stats-pool", 0L, 64 * 1024 * 1024, PoolGroup.SEARCH); NativeAllocatorPoolStats stats = allocator.stats(); assertNotNull(stats); - assertEquals(1024 * 1024 * 1024, stats.getRootLimitBytes()); - assertEquals(0, stats.getRootAllocatedBytes()); + assertEquals(-1, stats.getNativeAllocatedBytes()); + assertEquals(-1, stats.getNativeResidentBytes()); assertEquals(1, stats.getPools().size()); NativeAllocatorPoolStats.PoolStats poolStats = stats.getPools().get(0); assertEquals("stats-pool", poolStats.getName()); assertEquals(64 * 1024 * 1024, poolStats.getLimitBytes()); assertEquals(0, poolStats.getAllocatedBytes()); - // child_count is no longer rendered in stats; getPoolAllocator(...).getChildAllocators() - // is the runtime accessor for that detail if needed. } public void testStatsMultiplePools() { - allocator.getOrCreatePool("pool-a", 100 * 1024 * 1024); - allocator.getOrCreatePool("pool-b", 200 * 1024 * 1024); + allocator.getOrCreatePool("pool-a", 0L, 100 * 1024 * 1024, null); + allocator.getOrCreatePool("pool-b", 0L, 200 * 1024 * 1024, null); NativeAllocatorPoolStats stats = allocator.stats(); assertEquals(2, stats.getPools().size()); } public void testGetPoolNames() { - allocator.getOrCreatePool("alpha", 10 * 1024 * 1024); - allocator.getOrCreatePool("beta", 20 * 1024 * 1024); + allocator.getOrCreatePool("alpha", 0L, 10 * 1024 * 1024, null); + allocator.getOrCreatePool("beta", 0L, 20 * 1024 * 1024, null); assertTrue(allocator.getPoolNames().contains("alpha")); assertTrue(allocator.getPoolNames().contains("beta")); assertEquals(2, allocator.getPoolNames().size()); } - public void testRebalanceDistributesHeadroomToAllPools() { - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("active", 10 * 1024 * 1024, 100 * 1024 * 1024); - allocator.getOrCreatePool("idle", 10 * 1024 * 1024, 100 * 1024 * 1024); - - // Simulate activity: allocate in "active" pool. - BufferAllocator activeAlloc = allocator.getPoolAllocator("active"); - BufferAllocator child = activeAlloc.newChildAllocator("worker", 0, 100 * 1024 * 1024); - var buf = child.buffer(5 * 1024 * 1024); - - try { - allocator.rebalance(); - - // Active pool gets bonus headroom on top of its min. - long activeLimit = activeAlloc.getLimit(); - assertTrue("Active pool limit should exceed min after rebalance, got " + activeLimit, activeLimit > 10 * 1024 * 1024); - - // Idle pool also receives headroom: distributing to all pools (not just - // currently-active ones) avoids the dead-pool corner case where a pool - // with min = 0 starts at limit = 0 and can never make a first allocation. - // Idle pools that don't end up needing the headroom return it on the next - // tick once they remain at zero allocation. - long idleLimit = allocator.getPoolAllocator("idle").getLimit(); - assertTrue("Idle pool should also receive headroom, got " + idleLimit, idleLimit > 10 * 1024 * 1024); - } finally { - buf.close(); - child.close(); - } - } - - public void testRebalanceLetsZeroMinPoolAllocate() { - // Regression test: under the previous "active pools only" rebalance algorithm, - // a pool with min = 0 would start at limit = 0 (rebalancer-on path), be unable - // to allocate, never become "active", and so never receive a bonus — permanently - // dead. Distributing headroom across all pools fixes the chicken-and-egg. - allocator.setRebalanceInterval(60); - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("zero-min", 0L, 100 * 1024 * 1024); - try { - allocator.rebalance(); - BufferAllocator pool = allocator.getPoolAllocator("zero-min"); - assertTrue("Zero-min pool should receive headroom, got " + pool.getLimit(), pool.getLimit() > 0); - } finally { - allocator.setRebalanceInterval(0); - } - } - - public void testRebalanceNeverDropsBelowCurrentAllocation() { - allocator.setRootLimit(50 * 1024 * 1024); - allocator.getOrCreatePool("busy", 10 * 1024 * 1024); - - BufferAllocator pool = allocator.getPoolAllocator("busy"); - BufferAllocator child = pool.newChildAllocator("w", 0, 10 * 1024 * 1024); - var buf = child.buffer(8 * 1024 * 1024); // 8 MB allocated - - try { - allocator.rebalance(); - assertTrue("Limit should never drop below current allocation", pool.getLimit() >= pool.getAllocatedMemory()); - } finally { - buf.close(); - child.close(); - } - } - - public void testRebalanceWithNoPools() { - // Should not throw - allocator.rebalance(); - } - - public void testInitialLimitIsMaxWhenRebalancerDisabled() { - // Default tearDown allocator has rebalancer disabled (interval=0). - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("burst", 10 * 1024 * 1024, 100 * 1024 * 1024); - // With the rebalancer off, pools must start at their max so consumers can allocate - // immediately. Otherwise default-configured pools (min=0) would reject everything. + public void testInitialLimitIsMax() { + NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("burst", 10 * 1024 * 1024, 100 * 1024 * 1024, null); assertEquals(100 * 1024 * 1024, handle.limit()); } - public void testInitialLimitIsMinWhenRebalancerEnabled() { - // Enabling the rebalancer reverts to the original "guarantee + burst" semantics: - // pools start at min and grow via the next rebalance tick. - allocator.setRebalanceInterval(60); // any positive value enables the flag - NativeAllocator.PoolHandle handle = allocator.getOrCreatePool("guaranteed", 10 * 1024 * 1024, 100 * 1024 * 1024); - assertEquals(10 * 1024 * 1024, handle.limit()); - // Disable so subsequent tests aren't affected by the scheduled task. - allocator.setRebalanceInterval(0); - } - public void testCloseReleasesAllPools() { - allocator.getOrCreatePool("close-test", 10 * 1024 * 1024); + allocator.getOrCreatePool("close-test", 0L, 10 * 1024 * 1024, null); allocator.close(); assertTrue(allocator.getPoolNames().isEmpty()); // Recreate for tearDown - allocator = new ArrowNativeAllocator(1024L * 1024 * 1024); + allocator = new ArrowNativeAllocator(); } - public void testSetPoolMinRaisesLiveLimitWhenRebalancerOff() { - // setPoolMin must affect the live BufferAllocator immediately, not just the - // poolMins map. Otherwise it's a Dynamic setting that returns HTTP 200 and - // does nothing observable until the operator also enables the rebalancer. - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024); + public void testSetPoolMinRaisesLiveLimitWhenNeeded() { + allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024, null); BufferAllocator pool = allocator.getPoolAllocator("p"); long startLimit = pool.getLimit(); @@ -222,15 +130,45 @@ public void testSetPoolMinRaisesLiveLimitWhenRebalancerOff() { } public void testSetPoolMinDoesNotShrinkLiveLimit() { - // Dropping the min must not shrink an in-flight pool — the rebalancer is the - // only path that reduces limits, so a min change on its own should never - // reclaim capacity. - allocator.setRootLimit(100 * 1024 * 1024); - allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024); + allocator.getOrCreatePool("p", 0L, 100 * 1024 * 1024, null); BufferAllocator pool = allocator.getPoolAllocator("p"); long startLimit = pool.getLimit(); allocator.setPoolMin("p", 1L); assertEquals("dropping min must not shrink live limit", startLimit, pool.getLimit()); } + + public void testPoolGroupLimitListenerFiresWithCorrectSum() { + allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); + allocator.getOrCreatePool("write", 0L, 50 * 1024 * 1024, PoolGroup.INDEXING); + + java.util.concurrent.atomic.AtomicLong received = new java.util.concurrent.atomic.AtomicLong(-1); + allocator.addPoolGroupLimitListener(PoolGroup.INDEXING, received::set); + + // Change limits + allocator.setPoolEffectiveLimit("ingest", 60 * 1024 * 1024); + allocator.setPoolEffectiveLimit("write", 40 * 1024 * 1024); + + // Fire after all changes + allocator.firePoolGroupListeners(PoolGroup.INDEXING); + + // Listener should receive sum of effective limits: 60MB + 40MB = 100MB + assertEquals(100L * 1024 * 1024, received.get()); + } + + public void testPoolGroupLimitListenerNotFiredForOtherGroups() { + allocator.getOrCreatePool("flight", 0L, 50 * 1024 * 1024, PoolGroup.TRANSPORT); + allocator.getOrCreatePool("ingest", 0L, 80 * 1024 * 1024, PoolGroup.INDEXING); + + java.util.concurrent.atomic.AtomicLong transportReceived = new java.util.concurrent.atomic.AtomicLong(-1); + allocator.addPoolGroupLimitListener(PoolGroup.TRANSPORT, transportReceived::set); + + // Fire INDEXING group — should NOT trigger TRANSPORT listener + allocator.firePoolGroupListeners(PoolGroup.INDEXING); + assertEquals(-1, transportReceived.get()); + + // Fire TRANSPORT group — should trigger + allocator.firePoolGroupListeners(PoolGroup.TRANSPORT); + assertEquals(50L * 1024 * 1024, transportReceived.get()); + } } diff --git a/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java new file mode 100644 index 0000000000000..b92a477a67ea3 --- /dev/null +++ b/plugins/arrow-base/src/test/java/org/opensearch/arrow/allocator/NativeMemoryRebalancerTests.java @@ -0,0 +1,152 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.allocator; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.List; + +public class NativeMemoryRebalancerTests extends OpenSearchTestCase { + + private static final long MB = 1024L * 1024; + private static final long BUDGET = 100 * MB; + + private ArrowNativeAllocator allocator; + private NativeMemoryRebalancer rebalancer; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new ArrowNativeAllocator(); + allocator.setBudget(BUDGET); + rebalancer = new NativeMemoryRebalancer(allocator, () -> BUDGET, 0.75, 0.50, 0.10); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testPoolsStartAtMax() { + allocator.getOrCreatePool("a", 5 * MB, 40 * MB, null); + allocator.getOrCreatePool("b", 10 * MB, 50 * MB, null); + + assertEquals(40 * MB, allocator.getPoolAllocator("a").getLimit()); + assertEquals(50 * MB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testShrinksIdlePool() { + allocator.getOrCreatePool("idle", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 50 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (50 * MB * 0.8)); + + try { + long idleLimitBefore = allocator.getPoolAllocator("idle").getLimit(); + rebalancer.rebalance(); + long idleLimitAfter = allocator.getPoolAllocator("idle").getLimit(); + assertTrue("Idle pool should shrink, was " + idleLimitBefore + " now " + idleLimitAfter, idleLimitAfter < idleLimitBefore); + } finally { + buf.close(); + } + } + + public void testGrowsPressuredPoolAboveMax() { + allocator.getOrCreatePool("idle", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 20 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (20 * MB * 0.8)); + + try { + rebalancer.rebalance(); + long pressuredLimit = pressuredPool.getLimit(); + assertTrue("Pressured pool should grow above max (20MB), got " + pressuredLimit, pressuredLimit > 20 * MB); + } finally { + buf.close(); + } + } + + public void testNeverDropsBelowMin() { + allocator.getOrCreatePool("floored", 10 * MB, 50 * MB, null); + allocator.getOrCreatePool("pressured", 5 * MB, 50 * MB, null); + + BufferAllocator pressuredPool = allocator.getPoolAllocator("pressured"); + ArrowBuf buf = pressuredPool.buffer((long) (50 * MB * 0.8)); + + try { + for (int i = 0; i < 20; i++) { + rebalancer.rebalance(); + } + long flooredLimit = allocator.getPoolAllocator("floored").getLimit(); + assertTrue("Pool limit (" + flooredLimit + ") should not drop below min (10MB)", flooredLimit >= 10 * MB); + } finally { + buf.close(); + } + } + + public void testNoActionWhenNoPressure() { + allocator.getOrCreatePool("a", 5 * MB, 50 * MB, null); + allocator.getOrCreatePool("b", 5 * MB, 50 * MB, null); + + long limitA = allocator.getPoolAllocator("a").getLimit(); + long limitB = allocator.getPoolAllocator("b").getLimit(); + + rebalancer.rebalance(); + + assertEquals(limitA, allocator.getPoolAllocator("a").getLimit()); + assertEquals(limitB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testResetAllPoolsToMax() { + allocator.getOrCreatePool("a", 5 * MB, 40 * MB, null); + allocator.getOrCreatePool("b", 5 * MB, 50 * MB, null); + + BufferAllocator bPool = allocator.getPoolAllocator("b"); + ArrowBuf buf = bPool.buffer((long) (50 * MB * 0.8)); + rebalancer.rebalance(); + buf.close(); + + allocator.resetAllPoolsToMax(); + + assertEquals(40 * MB, allocator.getPoolAllocator("a").getLimit()); + assertEquals(50 * MB, allocator.getPoolAllocator("b").getLimit()); + } + + public void testSumLimitsNeverExceedsBudget() { + allocator.getOrCreatePool("p1", 5 * MB, 30 * MB, null); + allocator.getOrCreatePool("p2", 5 * MB, 30 * MB, null); + allocator.getOrCreatePool("p3", 5 * MB, 30 * MB, null); + + List bufs = new ArrayList<>(); + for (String name : new String[] { "p1", "p2", "p3" }) { + BufferAllocator pool = allocator.getPoolAllocator(name); + bufs.add(pool.buffer((long) (pool.getLimit() * 0.8))); + } + + try { + for (int i = 0; i < 10; i++) { + rebalancer.rebalance(); + } + + long sumLimits = 0; + for (String name : new String[] { "p1", "p2", "p3" }) { + sumLimits += allocator.getPoolAllocator(name).getLimit(); + } + assertTrue("Sum of limits (" + sumLimits + ") should not exceed budget (" + BUDGET + ")", sumLimits <= BUDGET); + } finally { + bufs.forEach(ArrowBuf::close); + } + } +} diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java index ff7f1359df089..a32a0ffe8bc58 100644 --- a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/BackpressureProducerIT.java @@ -122,9 +122,10 @@ protected Settings nodeSettings(int nodeOrdinal) { return Settings.builder() .put(super.nodeSettings(nodeOrdinal)) .put("node.native_memory.limit", "512mb") - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, 256 * MB) .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, FLIGHT_POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 8 * MB) .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 16 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 8 * MB) .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 16 * MB) .put(ServerConfig.FLIGHT_OUTBOUND_BUFFER_THRESHOLD.getKey(), new ByteSizeValue(GRPC_THRESHOLD_BYTES, ByteSizeUnit.BYTES)) .build(); diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java index 5633cfa429c5d..d849458c0b048 100644 --- a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeAllocatorBoundaryIT.java @@ -12,7 +12,6 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.OutOfMemoryException; -import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse; import org.opensearch.arrow.allocator.ArrowBasePlugin; import org.opensearch.arrow.allocator.ArrowNativeAllocator; import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; @@ -31,27 +30,17 @@ * *

Boots a single-node cluster with tight memory settings, then exercises * the actual Arrow allocation path to verify that the framework's - * configured caps are enforced at allocation time (not just at config-parse - * time). Complements unit-level tests in {@code ArrowBasePluginTests} by - * verifying that the production wiring (Guice -> ArrowNativeAllocator -> - * Arrow's RootAllocator chain) honors the caps end-to-end. - * - *

Each test sets explicit byte limits and allocates real - * {@link org.apache.arrow.memory.ArrowBuf} buffers, asserting either - * successful allocation or {@link OutOfMemoryException} based on whether - * the request fits within the configured cap. + * configured caps are enforced at allocation time. */ @ThreadLeakScope(ThreadLeakScope.Scope.NONE) @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) public class NativeAllocatorBoundaryIT extends OpenSearchIntegTestCase { - /** 1 MiB. Chosen small enough that tests run fast but large enough that - * Arrow's internal accounting doesn't round it away. */ + /** 1 MiB. */ private static final long MB = 1024L * 1024; - /** Cap large enough for the framework's own bookkeeping but small enough - * to trigger OOM well before exhausting host memory. */ - private static final long ROOT_CAP_BYTES = 16 * MB; + /** Per-pool cap for tests. */ + private static final long POOL_CAP_BYTES = 16 * MB; @Override protected Collection> nodePlugins() { @@ -60,127 +49,83 @@ protected Collection> nodePlugins() { @Override protected Settings nodeSettings(int nodeOrdinal) { - // Set node.native_memory.limit explicitly so framework defaults derive - // from a known value rather than the (machine-dependent) ram-heap default. - // ROOT_LIMIT and pool maxes are then overridden per-test via cluster - // settings PUT or directly via this node-settings layer. return Settings.builder() .put(super.nodeSettings(nodeOrdinal)) .put("node.native_memory.limit", "256mb") - // Tight root cap: 16 MiB total Arrow framework budget. - .put(NativeAllocatorPoolConfig.SETTING_ROOT_LIMIT, ROOT_CAP_BYTES) - // Per-pool maxes set generously so per-pool caps don't trip - // before root.limit. Tests targeting per-pool caps override below. - .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, ROOT_CAP_BYTES) - .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, ROOT_CAP_BYTES) - .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, ROOT_CAP_BYTES) + .put("native.allocator.rebalancer.enabled", false) + // Per-pool maxes set to a known value for testing. + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, POOL_CAP_BYTES) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, POOL_CAP_BYTES) .build(); } /** - * Verifies that {@code parquet.native.pool.query.max} caps allocations - * through the QUERY pool: a buffer request exceeding the per-pool cap - * throws {@link OutOfMemoryException} even when root has headroom. + * Verifies that per-pool max caps allocations through the QUERY pool. */ public void testPoolMaxRejectsAllocationsBeyondCap() { - // Tighten QUERY pool to 4 MiB while leaving root at 16 MiB. - long poolCap = 4 * MB; - ClusterUpdateSettingsResponse resp = client().admin() - .cluster() - .prepareUpdateSettings() - .setTransientSettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, poolCap)) - .get(); - assertTrue("PUT to query.max must succeed", resp.isAcknowledged()); - ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - assertThat("pool's live limit reflects the PUT", queryPool.getLimit(), is(poolCap)); + assertThat("pool's live limit reflects the configured max", queryPool.getLimit(), is(POOL_CAP_BYTES)); // Sub-cap allocation succeeds. try (var withinCap = queryPool.buffer(2 * MB)) { assertThat(queryPool.getAllocatedMemory(), greaterThanOrEqualTo(2 * MB)); } - // Cap+1 allocation fails — Arrow's parent-cap check at allocateBytes - // walks queryPool's allocationLimit and rejects. - expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(8 * MB)); + // Cap+1 allocation fails. + expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(POOL_CAP_BYTES + MB)); } /** - * Verifies that {@code native.allocator.root.limit} caps allocations - * across all pools combined: when the sum of in-flight pool allocations - * approaches the root cap, the next allocation is rejected at the root - * level even if each individual pool's max would allow it. + * Verifies that per-pool limits cap allocations: when one pool is full, + * allocations through it fail even if other pools have headroom. */ - public void testRootLimitRejectsAllocationsBeyondCap() { + public void testPoolLimitRejectsAllocationsBeyondCap() { ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator flightPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - BufferAllocator root = allocator.getRootAllocator(); - - // Sanity-check setup: confirm the live limits match nodeSettings. - // If these fail, the test setup is wrong and the body's expectations are - // meaningless — surface the misconfiguration instead of misleading failures below. - assertThat("root.limit must match nodeSettings", root.getLimit(), is(ROOT_CAP_BYTES)); - assertThat("flight.max must match nodeSettings", flightPool.getLimit(), is(ROOT_CAP_BYTES)); - assertThat("query.max must match nodeSettings", queryPool.getLimit(), is(ROOT_CAP_BYTES)); - - // Hold 8 MiB through the FLIGHT pool. With root at 16 MiB this leaves 8 MiB - // headroom across the root. (Power-of-2 sizes avoid Arrow's chunked-allocation - // rounding surprises; e.g. a 12 MiB request actually consumes 16 MiB.) + + assertThat("flight.max must match nodeSettings", flightPool.getLimit(), is(POOL_CAP_BYTES)); + assertThat("query.max must match nodeSettings", queryPool.getLimit(), is(POOL_CAP_BYTES)); + + // Hold 8 MiB through the FLIGHT pool. try (var flightHold = flightPool.buffer(8 * MB)) { assertThat("FLIGHT pool reflects 8MB allocation", flightPool.getAllocatedMemory(), is(8L * MB)); - assertThat("root reflects 8MB allocation", root.getAllocatedMemory(), is(8L * MB)); - // A 4 MiB allocation through QUERY succeeds (within remaining root headroom). + // A 4 MiB allocation through QUERY succeeds (within its own pool cap). try (var queryFit = queryPool.buffer(4 * MB)) { - assertThat(allocator.getRootAllocator().getAllocatedMemory(), is(12L * MB)); + assertThat(queryPool.getAllocatedMemory(), greaterThanOrEqualTo(4 * MB)); } - // An 8 MiB allocation through QUERY would push the root past the 16 MiB cap - // (8 MiB FLIGHT + 8 MiB QUERY). Arrow's parent-cap check at allocateBytes - // walks queryPool -> root and rejects with OOM, even though QUERY's own - // (16 MiB) max would individually allow it. - expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(16 * MB)); + // An allocation exceeding the QUERY pool's own cap fails. + expectThrows(OutOfMemoryException.class, () -> queryPool.buffer(POOL_CAP_BYTES + MB)); } } /** - * Verifies that a dynamic PUT to a pool's max takes effect on - * subsequent allocations through descendants of that pool. This is the - * behavior the deleted {@code NativeAllocatorListener} SPI was emulating; - * it is now provided natively by Arrow's parent-cap check at allocateBytes. + * Verifies that setPoolLimit dynamically adjusts the pool cap and + * subsequent allocations respect the new limit. */ - public void testDynamicPoolResizeAffectsInFlightAllocations() { + public void testSetPoolLimitAffectsInFlightAllocations() { ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); BufferAllocator queryPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); - // Step 1: create a child allocator at Long.MAX_VALUE — the AnalyticsSearchService / - // DefaultPlanExecutor pattern. The child intentionally has no own-cap; it relies - // on the parent pool's allocationLimit at allocation time. try (BufferAllocator child = queryPool.newChildAllocator("boundary-it-child", 0, Long.MAX_VALUE)) { - // Step 2: a small buffer through the child succeeds with the initial pool max. + // A small buffer through the child succeeds with the initial pool max. try (var buf = child.buffer(2 * MB)) { assertThat(child.getAllocatedMemory(), greaterThanOrEqualTo(2 * MB)); } - // Step 3: PUT a tighter pool max via cluster settings. + // Programmatically tighten the pool limit. long newPoolCap = 1 * MB; - ClusterUpdateSettingsResponse resp = client().admin() - .cluster() - .prepareUpdateSettings() - .setTransientSettings(Settings.builder().put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, newPoolCap)) - .get(); - assertTrue("PUT to query.max must succeed", resp.isAcknowledged()); - assertThat("pool's own limit reflects the PUT", queryPool.getLimit(), is(newPoolCap)); - assertThat("child's own limit is intentionally uncapped", child.getLimit(), is(Long.MAX_VALUE)); - - // Step 4: an allocation that fit before the resize now exceeds the parent cap. - // Arrow's parent-cap check at allocateBytes walks queryPool.allocationLimit - // and rejects — no listener machinery needed. + allocator.setPoolLimit(NativeAllocatorPoolConfig.POOL_QUERY, newPoolCap); + assertThat("pool's own limit reflects the update", queryPool.getLimit(), is(newPoolCap)); + + // An allocation that fit before the resize now exceeds the parent cap. expectThrows(OutOfMemoryException.class, () -> child.buffer(2 * MB)); - // Step 5: an allocation under the new cap still succeeds. + // An allocation under the new cap still succeeds. try (var smallBuf = child.buffer(512 * 1024)) { assertThat(child.getAllocatedMemory(), greaterThanOrEqualTo(512L * 1024)); } diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java new file mode 100644 index 0000000000000..e28e373fe69a9 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/NativeMemoryRebalancerIT.java @@ -0,0 +1,100 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.arrow.allocator.ArrowNativeAllocator; +import org.opensearch.arrow.spi.NativeAllocatorPoolConfig; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Collection; +import java.util.List; + +import static org.hamcrest.Matchers.lessThanOrEqualTo; + +/** + * Integration test for the NativeMemoryRebalancer. + * + *

Boots a single-node cluster with the rebalancer enabled and verifies that + * pools start at their max and the rebalancer shrinks idle pools / grows pressured ones. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) +public class NativeMemoryRebalancerIT extends OpenSearchIntegTestCase { + + private static final long MB = 1024L * 1024; + + @Override + protected Collection> nodePlugins() { + return List.of(ArrowBasePlugin.class); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put("native.allocator.rebalancer.enabled", true) + .put("native.allocator.rebalance.interval_seconds", 1) + .put("node.native_memory.limit", "1gb") + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_FLIGHT_MAX, 200 * MB) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_INGEST_MAX, 200 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MIN, 5 * MB) + .put(NativeAllocatorPoolConfig.SETTING_QUERY_MAX, 200 * MB) + .build(); + } + + /** + * Verifies that pools start at their max (before rebalancer shrinks them). + * Uses a long rebalancer interval to avoid race conditions. + */ + public void testPoolsStartAtMax() { + ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); + BufferAllocator ingestPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST); + + // The rebalancer may have already run (1s interval), so the pool may have shrunk. + // Verify it's at least at min and the configured max is correct. + long max = allocator.getPoolMax(NativeAllocatorPoolConfig.POOL_INGEST); + assertEquals("Ingest pool max should be configured at 200MB", 200 * MB, max); + // Pool limit should be between min and max (rebalancer may have shrunk it) + long limit = ingestPool.getLimit(); + assertThat("Ingest pool limit should be >= min", limit, org.hamcrest.Matchers.greaterThanOrEqualTo(5 * MB)); + assertThat("Ingest pool limit should be <= max", limit, lessThanOrEqualTo(200 * MB)); + } + + /** + * Verifies that an idle pool shrinks after rebalancer ticks when another pool is pressured. + */ + public void testIdlePoolShrinksWhenOtherPressured() throws Exception { + ArrowNativeAllocator allocator = internalCluster().getInstance(ArrowNativeAllocator.class); + BufferAllocator ingestPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_INGEST); + BufferAllocator flightPool = allocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_FLIGHT); + + // Allocate > 75% of ingest pool to create pressure + long toAllocate = (long) (ingestPool.getLimit() * 0.8); + ArrowBuf buf = ingestPool.buffer(toAllocate); + + try { + // Flight pool is idle — wait for rebalancer to shrink it + assertBusy(() -> { + long flightLimit = flightPool.getLimit(); + assertThat("Flight pool should shrink when idle", flightLimit, org.hamcrest.Matchers.lessThan(200 * MB)); + }); + } finally { + buf.close(); + } + } +} diff --git a/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java new file mode 100644 index 0000000000000..fd0e165103e86 --- /dev/null +++ b/plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/UnifiedNativeMemoryStatsIT.java @@ -0,0 +1,95 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.arrow.flight; + +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; + +import org.opensearch.action.admin.cluster.node.stats.NodesStatsRequest; +import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.common.settings.Settings; +import org.opensearch.plugin.stats.NativeAllocatorPoolStats; +import org.opensearch.plugins.Plugin; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Collection; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasItems; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.notNullValue; + +/** + * Integration test verifying the unified native memory stats endpoint. + * Boots a single-node cluster with ArrowBasePlugin and confirms that + * all registered pools (Arrow + virtual) appear in _nodes/stats/native_memory. + */ +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, minNumDataNodes = 1, maxNumDataNodes = 1) +public class UnifiedNativeMemoryStatsIT extends OpenSearchIntegTestCase { + + @Override + protected Collection> nodePlugins() { + return List.of(ArrowBasePlugin.class); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder().put(super.nodeSettings(nodeOrdinal)).put("node.native_memory.limit", "1gb").build(); + } + + /** + * Verifies that the Arrow pools (flight, ingest, query) are visible in + * _nodes/stats/native_memory with correct structure. + */ + public void testArrowPoolsVisibleInStats() { + NodesStatsResponse response = client().admin() + .cluster() + .prepareNodesStats() + .addMetric(NodesStatsRequest.Metric.NATIVE_MEMORY.metricName()) + .get(); + + assertThat(response.getNodes().isEmpty(), is(false)); + NativeAllocatorPoolStats stats = response.getNodes().get(0).getNativeAllocatorStats(); + assertThat("native_memory stats should be present", stats, notNullValue()); + + // Dump the stats for debugging + StringBuilder sb = new StringBuilder(); + sb.append("nativeAllocated=").append(stats.getNativeAllocatedBytes()); + sb.append(", nativeResident=").append(stats.getNativeResidentBytes()); + sb.append(", pools=["); + for (NativeAllocatorPoolStats.PoolStats p : stats.getPools()) { + sb.append(p.getName()) + .append("(alloc=") + .append(p.getAllocatedBytes()) + .append(",peak=") + .append(p.getPeakBytes()) + .append(",limit=") + .append(p.getLimitBytes()) + .append(") "); + } + sb.append("]"); + logger.info("=== NATIVE_MEMORY STATS: {} ===", sb); + + // All Arrow pools should be present + Set poolNames = stats.getPools().stream().map(NativeAllocatorPoolStats.PoolStats::getName).collect(Collectors.toSet()); + assertThat(poolNames, hasItems("flight", "ingest", "query")); + + // Each pool should have limit > 0 (derived from 1gb native_memory.limit) + for (NativeAllocatorPoolStats.PoolStats pool : stats.getPools()) { + assertThat("Pool '" + pool.getName() + "' should have limit > 0", pool.getLimitBytes(), greaterThan(0L)); + assertThat("Pool '" + pool.getName() + "' allocated should be >= 0", pool.getAllocatedBytes(), greaterThanOrEqualTo(0L)); + } + } + +} diff --git a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java index f5eacc771be92..add8e0a2bbfee 100644 --- a/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java +++ b/plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/transport/FlightTransportTestBase.java @@ -91,8 +91,8 @@ public void setUp() throws Exception { // FlightTransport sources its allocator from the framework's FLIGHT pool. Construct one // here so the test has a usable allocator; tearDown closes it. - nativeAllocator = new ArrowNativeAllocator(Long.MAX_VALUE); - nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_FLIGHT, 0L, Long.MAX_VALUE); + nativeAllocator = new ArrowNativeAllocator(); + nativeAllocator.getOrCreatePool(NativeAllocatorPoolConfig.POOL_FLIGHT, 0L, Long.MAX_VALUE, null); flightTransport = new FlightTransport( settings, diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 deleted file mode 100644 index 02ab255e34d5e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-api-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -c4ee83d77005567852a72e08b945ebb023be1daa \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0f145e077247a --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-api-1.63.0.jar.sha1 @@ -0,0 +1 @@ +39c923c0f236417ec8c4e2091f9e3032b4b6fb91 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 deleted file mode 100644 index 88802c8009d0e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.62.0-alpha.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -91f3bcf6b93261cbaf32dd156e0007aa5fa5b25a \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 new file mode 100644 index 0000000000000..aa0966ff936aa --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-api-incubator-1.63.0-alpha.jar.sha1 @@ -0,0 +1 @@ +2a88356ec37eb66666dc6798e4221b26af74e24d \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 deleted file mode 100644 index db25f474db864..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -e6468bd64a94429b68761f7c13e143c3fdfaafc7 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..5c57ec9f8ba92 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +450f8d552f33d51b19457d1336d8f7bdaec35e01 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 deleted file mode 100644 index 4608aebd30520..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-context-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -365cee4d1f365e4d4a05654742b50aa436c2dd8e \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0b6d3c50f42e3 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-context-1.63.0.jar.sha1 @@ -0,0 +1 @@ +4af6d513cedbf78dafe104f82f8dce7c7a205f07 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 deleted file mode 100644 index 9114878e88cef..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -2dafa6ae65cbf1aa321cd644d200f3ff8465284d \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..20fd2fafeff60 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +75e2658228eb885770b894f71e8d145d46095fdf \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 deleted file mode 100644 index 25565231ce2ff..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -79ec5f1f23e00da7a8c8a30136cfbfaf9aa38f93 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..aaaef68da4186 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-logging-1.63.0.jar.sha1 @@ -0,0 +1 @@ +0a5f203ebd35f1ff54171a01335d1a47c3600523 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 deleted file mode 100644 index 1d0b19032d4ef..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -8e4cb9199ac868332a1213ca27408a18905ba369 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..6b5d64d4dfe72 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-1.63.0.jar.sha1 @@ -0,0 +1 @@ +7f02be0f4f2c44476e957bce9d7f7d3e0a1a7ae4 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 deleted file mode 100644 index 03a22b1f68946..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -6f269df0e3f065fbd2e590458e7b2107cde2a106 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..e1adbf694257b --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-otlp-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +d6a9a75d3e4b9f69a67c3a333162a3b0783178c0 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 deleted file mode 100644 index 82a20ade44ef5..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -19b5e023db9037a38fe2531afb6e44456e963fba \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..71c90cfe4405d --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-exporter-sender-okhttp-1.63.0.jar.sha1 @@ -0,0 +1 @@ +2045d90ed21954e0f61e900d168034e17711c7e9 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 deleted file mode 100644 index 99ef9c2e1d260..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -06fa52c4641322b14b8bd515eb048bb9b1365d0c \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..5cc2625880df0 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-1.63.0.jar.sha1 @@ -0,0 +1 @@ +1d522dcf3fb4903c7f95835c3f6f631cf1c3824b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 deleted file mode 100644 index 9828f6e0985cf..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -b6742282daab8e13598b78a83ddfa54f10b5752b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..c4a8a431030c5 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-common-1.63.0.jar.sha1 @@ -0,0 +1 @@ +db2ec1f09c307adcf015bb8a87dd687eefa32e77 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 deleted file mode 100644 index b008d4cb9a80e..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -f242422084100da0bd3a5f6f2bcf364aaf4d2c53 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..dfdc2f6853b85 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-logs-1.63.0.jar.sha1 @@ -0,0 +1 @@ +ef1e93dc6e05d2b9191d7fb9449fea8965086acb \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 deleted file mode 100644 index a845283b6a3d4..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -5838371075930a4a15f7f61240b4b64cb3e924d8 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..2db28d76c8a47 --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-metrics-1.63.0.jar.sha1 @@ -0,0 +1 @@ +decbaf0de6dbe8156aa426ddc85b5b38354a236b \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 deleted file mode 100644 index 9d3ce157565d3..0000000000000 --- a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.62.0.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -7a337d2f887b151d27e734d1c221eb51b1c5b734 \ No newline at end of file diff --git a/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 new file mode 100644 index 0000000000000..0ee645084e84b --- /dev/null +++ b/plugins/telemetry-otel/licenses/opentelemetry-sdk-trace-1.63.0.jar.sha1 @@ -0,0 +1 @@ +5c948707623a0dbd79105eb29461f42bfe0db11b \ No newline at end of file diff --git a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java index 0efcb47a7d927..88f8bf7ff67be 100644 --- a/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java +++ b/sandbox/libs/analytics-api/src/main/java/org/opensearch/analytics/exec/profile/TaskProfile.java @@ -12,17 +12,23 @@ import org.opensearch.core.xcontent.XContentBuilder; import java.io.IOException; +import java.util.Map; /** - * Per-task profile snapshot. Captures target node, terminal state and - * wall-clock elapsed time. A "task" is one dispatch unit within a stage - * (one shard for SOURCE, one partition for HASH_PARTITIONED, one total for COORDINATOR). + * Per-task profile snapshot. Captures target node, terminal state, + * wall-clock elapsed time, and optional data-node execution metrics. * * @param node target node and shard the task ran on, or "(unknown)" if dispatch never happened * @param state terminal state — CREATED if the task was never dispatched * @param elapsedMs wall-clock time from dispatch to terminal, or 0 if never dispatched + * @param dataNodeMetrics execution metrics from the data node (DataFusion operator timings), or null if not profiled */ -public record TaskProfile(String node, String state, long elapsedMs) implements ToXContentObject { +public record TaskProfile(String node, String state, long elapsedMs, Map dataNodeMetrics) implements ToXContentObject { + + /** Convenience constructor without data node metrics. */ + public TaskProfile(String node, String state, long elapsedMs) { + this(node, state, elapsedMs, null); + } @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { @@ -30,6 +36,13 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field("node", node); builder.field("state", state); builder.field("elapsed_ms", elapsedMs); + if (dataNodeMetrics != null && dataNodeMetrics.isEmpty() == false) { + builder.startObject("data_node_metrics"); + for (Map.Entry entry : dataNodeMetrics.entrySet()) { + builder.field(entry.getKey(), entry.getValue()); + } + builder.endObject(); + } builder.endObject(); return builder; } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java index 6b6f27440a470..6fbacfa86e8b9 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java @@ -38,6 +38,10 @@ public enum ScalarFunction { LESS_THAN_OR_EQUAL(Category.COMPARISON, SqlKind.LESS_THAN_OR_EQUAL), IS_NULL(Category.COMPARISON, SqlKind.IS_NULL), IS_NOT_NULL(Category.COMPARISON, SqlKind.IS_NOT_NULL), + IS_TRUE(Category.COMPARISON, SqlKind.IS_TRUE), + IS_FALSE(Category.COMPARISON, SqlKind.IS_FALSE), + IS_NOT_TRUE(Category.COMPARISON, SqlKind.IS_NOT_TRUE), + IS_NOT_FALSE(Category.COMPARISON, SqlKind.IS_NOT_FALSE), IN(Category.COMPARISON, SqlKind.IN), LIKE(Category.COMPARISON, SqlKind.LIKE), PREFIX(Category.COMPARISON, SqlKind.OTHER_FUNCTION), diff --git a/sandbox/libs/dataformat-native/rust/common/src/lib.rs b/sandbox/libs/dataformat-native/rust/common/src/lib.rs index 0f4b8c132407f..c44fa871c4fb3 100644 --- a/sandbox/libs/dataformat-native/rust/common/src/lib.rs +++ b/sandbox/libs/dataformat-native/rust/common/src/lib.rs @@ -11,6 +11,7 @@ pub mod error; pub mod logger; pub mod allocator; +pub mod memory_pool; // Re-export the proc macro so plugins use `#[native_bridge_common::ffm_safe]` pub use native_bridge_macros::ffm_safe; diff --git a/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs b/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs new file mode 100644 index 0000000000000..7fc5c7c83a20c --- /dev/null +++ b/sandbox/libs/dataformat-native/rust/common/src/memory_pool.rs @@ -0,0 +1,370 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +//! Memory pool for tracking native memory usage across write and merge operations. +//! +//! Provides an atomic counter with a configurable limit. Operations that allocate +//! significant memory call `try_grow` before allocating and `shrink` after freeing. +//! The pool rejects allocations that would exceed the configured limit. +//! +//! `MemoryReservation` is an RAII handle that automatically returns memory to the +//! pool on drop, preventing leaks even on error paths. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::Duration; +use std::fmt; + +/// Default timeout for blocking wait (300 seconds). +pub const DEFAULT_WAIT_TIMEOUT: Duration = Duration::from_secs(300); + +/// Merge operations can wait longer (600 seconds). +pub const MERGE_WAIT_TIMEOUT: Duration = Duration::from_secs(600); + +/// Controls whether an allocation blocks or rejects immediately. +#[derive(Debug, Clone)] +pub enum PoolBehavior { + /// Block until memory is available, up to the given timeout. + Wait(Duration), + /// Fail immediately if pool is full. + Reject, +} + +/// Error returned when a pool cannot satisfy an allocation request. +#[derive(Debug, Clone)] +pub struct PoolExhausted { + pub pool_name: &'static str, + pub requested: usize, + pub used: usize, + pub limit: usize, +} + +impl fmt::Display for PoolExhausted { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "[{}] memory limit exceeded: requested {} bytes, used {}, limit {}", + self.pool_name, self.requested, self.used, self.limit + ) + } +} + +impl std::error::Error for PoolExhausted {} + +/// Error returned when wait_and_grow times out. +#[derive(Debug, Clone)] +pub struct PoolTimeout { + pub pool_name: &'static str, + pub requested: usize, + pub used: usize, + pub limit: usize, + pub waited: Duration, +} + +impl fmt::Display for PoolTimeout { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "[{}] timed out waiting for {} bytes after {:?} (used: {}, limit: {})", + self.pool_name, self.requested, self.waited, self.used, self.limit + ) + } +} + +impl std::error::Error for PoolTimeout {} + +/// A node-level memory pool backed by an atomic counter with blocking wait support. +pub struct MemoryPool { + name: &'static str, + used: AtomicUsize, + limit: AtomicUsize, + peak: AtomicUsize, + notify: Condvar, + notify_lock: Mutex<()>, +} + +impl fmt::Debug for MemoryPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MemoryPool") + .field("name", &self.name) + .field("used", &self.used.load(Ordering::Relaxed)) + .field("limit", &self.limit.load(Ordering::Relaxed)) + .field("peak", &self.peak.load(Ordering::Relaxed)) + .finish() + } +} + +impl MemoryPool { + /// Create a new pool. `limit = 0` means unlimited. + pub fn new(name: &'static str, limit: usize) -> Self { + Self { + name, + used: AtomicUsize::new(0), + limit: AtomicUsize::new(limit), + peak: AtomicUsize::new(0), + notify: Condvar::new(), + notify_lock: Mutex::new(()), + } + } + + /// Attempt to reserve `bytes`. Returns error if it would exceed the limit. + pub fn try_grow(&self, bytes: usize) -> Result<(), PoolExhausted> { + if bytes == 0 { + return Ok(()); + } + let limit = self.limit.load(Ordering::Relaxed); + let result = self.used.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| { + let new_used = used.checked_add(bytes)?; + if limit > 0 && new_used > limit { + None + } else { + Some(new_used) + } + }); + + match result { + Ok(old) => { + self.peak.fetch_max(old + bytes, Ordering::Relaxed); + Ok(()) + } + Err(_) => Err(PoolExhausted { + pool_name: self.name, + requested: bytes, + used: self.used.load(Ordering::Relaxed), + limit, + }), + } + } + + /// Blocks until `bytes` can be reserved, or timeout expires. + pub fn wait_and_grow(&self, bytes: usize, timeout: Duration) -> Result<(), PoolTimeout> { + if bytes == 0 { + return Ok(()); + } + if self.try_grow(bytes).is_ok() { + return Ok(()); + } + + let start = std::time::Instant::now(); + loop { + let elapsed = start.elapsed(); + if elapsed >= timeout { + let used = self.used.load(Ordering::Relaxed); + let limit = self.limit.load(Ordering::Relaxed); + return Err(PoolTimeout { + pool_name: self.name, + requested: bytes, + used, + limit, + waited: elapsed, + }); + } + + let remaining = timeout - elapsed; + let guard = self.notify_lock.lock().unwrap(); + let _ = self.notify.wait_timeout(guard, remaining.min(Duration::from_secs(1))).unwrap(); + + if self.try_grow(bytes).is_ok() { + return Ok(()); + } + } + } + + /// Infallible grow — use when the allocation has already happened. + pub fn grow(&self, bytes: usize) { + if bytes == 0 { + return; + } + let new_used = self.used.fetch_add(bytes, Ordering::Relaxed) + bytes; + self.peak.fetch_max(new_used, Ordering::Relaxed); + } + + /// Release `bytes` back to the pool. Notifies any waiting threads. + pub fn shrink(&self, bytes: usize) { + if bytes == 0 { + return; + } + self.used + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(bytes)) + }) + .unwrap(); + self.notify.notify_all(); + } + + pub fn used(&self) -> usize { + self.used.load(Ordering::Relaxed) + } + + pub fn peak(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + pub fn limit(&self) -> usize { + self.limit.load(Ordering::Relaxed) + } + + pub fn name(&self) -> &'static str { + self.name + } + + /// Atomically update the limit. Called by the Java rebalancer. + pub fn set_limit(&self, new_limit: usize) { + self.limit.store(new_limit, Ordering::Release); + // Wake waiters — new limit might allow blocked allocations + self.notify.notify_all(); + } +} + +/// RAII handle that tracks a portion of memory reserved from a [`MemoryPool`]. +/// Automatically releases all held memory on drop. +pub struct MemoryReservation { + pool: Arc, + consumer: &'static str, + size: usize, + behavior: PoolBehavior, +} + +impl MemoryReservation { + pub fn new(pool: &Arc, consumer: &'static str, behavior: PoolBehavior) -> Self { + Self { + pool: Arc::clone(pool), + consumer, + size: 0, + behavior, + } + } + + /// Grow based on the reservation's behavior: block (Wait) or reject (Reject). + pub fn request(&mut self, bytes: usize) -> Result<(), Box> { + match &self.behavior { + PoolBehavior::Reject => { + self.pool.try_grow(bytes)?; + self.size += bytes; + Ok(()) + } + PoolBehavior::Wait(timeout) => { + self.pool.wait_and_grow(bytes, *timeout)?; + self.size += bytes; + Ok(()) + } + } + } + + /// Infallible grow. + pub fn grow(&mut self, bytes: usize) { + self.pool.grow(bytes); + self.size += bytes; + } + + /// Release `bytes` from this reservation. + pub fn shrink(&mut self, bytes: usize) { + let actual = bytes.min(self.size); + self.pool.shrink(actual); + self.size -= actual; + } + + /// Release all memory back to the pool. + pub fn free(&mut self) -> usize { + let s = self.size; + if s > 0 { + self.pool.shrink(s); + self.size = 0; + } + s + } + + pub fn size(&self) -> usize { + self.size + } + + pub fn consumer(&self) -> &'static str { + self.consumer + } + + /// Create a sibling reservation from the same pool with a different consumer name. + pub fn child(&self, consumer: &'static str) -> Self { + Self { + pool: Arc::clone(&self.pool), + consumer, + size: 0, + behavior: self.behavior.clone(), + } + } +} + +impl Drop for MemoryReservation { + fn drop(&mut self) { + if self.size > 0 { + self.pool.shrink(self.size); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_try_grow_within_limit() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(512).is_ok()); + assert_eq!(res.size(), 512); + assert_eq!(pool.used(), 512); + } + + #[test] + fn test_try_grow_exceeds_limit() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(2048).is_err()); + assert_eq!(res.size(), 0); + assert_eq!(pool.used(), 0); + } + + #[test] + fn test_drop_releases_memory() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + { + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + res.request(500).unwrap(); + assert_eq!(pool.used(), 500); + } + assert_eq!(pool.used(), 0); + } + + #[test] + fn test_set_limit_allows_growth() { + let pool = Arc::new(MemoryPool::new("test", 100)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + assert!(res.request(200).is_err()); + pool.set_limit(500); + assert!(res.request(200).is_ok()); + } + + #[test] + fn test_peak_tracking() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let mut res = MemoryReservation::new(&pool, "writer", PoolBehavior::Reject); + res.request(800).unwrap(); + res.shrink(500); + assert_eq!(pool.peak(), 800); + assert_eq!(pool.used(), 300); + } + + #[test] + fn test_child_reservation() { + let pool = Arc::new(MemoryPool::new("test", 1024)); + let res = MemoryReservation::new(&pool, "parent", PoolBehavior::Reject); + let mut child = res.child("child"); + child.request(100).unwrap(); + assert_eq!(child.consumer(), "child"); + assert_eq!(pool.used(), 100); + } +} diff --git a/sandbox/libs/plugin-stats-spi/build.gradle b/sandbox/libs/plugin-stats-spi/build.gradle index 405dd7437a633..c12544b4caab6 100644 --- a/sandbox/libs/plugin-stats-spi/build.gradle +++ b/sandbox/libs/plugin-stats-spi/build.gradle @@ -14,9 +14,7 @@ java { sourceCompatibility = JavaVersion.toVersion(25); targetCompatibility = JavaVersion.toVersion(25) } -// no tests for now -testingConventions.enabled = false - dependencies { compileOnly project(':server') + testImplementation project(':test:framework') } diff --git a/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java b/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java index 249eaa65851e6..1ff03dc361da5 100644 --- a/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java +++ b/sandbox/libs/plugin-stats-spi/src/main/java/org/opensearch/plugin/stats/StatsRecorder.java @@ -72,6 +72,22 @@ public static void recordTimeMillis(IORunnable work, LongConsumer timeRecorder) } } + /** + * Value-returning variant of {@link #recordTimeMillis(IORunnable, LongConsumer)}. + * Runs {@code work}, returns its result, and reports elapsed wall-clock millis to + * {@code timeRecorder}. Time is reported even if {@code work} throws. + */ + public static T recordTimeMillis(IOSupplier work, LongConsumer timeRecorder) throws IOException { + Objects.requireNonNull(work, "work"); + Objects.requireNonNull(timeRecorder, "timeRecorder"); + long start = System.nanoTime(); + try { + return work.get(); + } finally { + timeRecorder.accept(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)); + } + } + /** * Runs {@code work} and reports time + outcome: *

    diff --git a/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java b/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java new file mode 100644 index 0000000000000..3284753213d79 --- /dev/null +++ b/sandbox/libs/plugin-stats-spi/src/test/java/org/opensearch/plugin/stats/StatsRecorderTests.java @@ -0,0 +1,247 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.plugin.stats; + +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Unit tests for {@link StatsRecorder}. + * + *

    Covers all four public methods and their contracts: + *

      + *
    • time is always recorded (success and failure paths)
      • + *
    • {@code onSuccess}/{@code onFailure} are mutually exclusive
      • + *
    • the original exception propagates unchanged
      • + *
    • return values pass through
      • + *
    • null arguments are rejected
      • + *
    + */ +public class StatsRecorderTests extends OpenSearchTestCase { + + // ────────────────────────────────────────────────────────────────────── + // recordTimeMillis(IORunnable, ...) + // ────────────────────────────────────────────────────────────────────── + + public void testRecordTimeMillisVoidRunsWorkAndRecordsTime() throws IOException { + AtomicBoolean ran = new AtomicBoolean(false); + AtomicLong recorded = new AtomicLong(-1); + + StatsRecorder.recordTimeMillis(() -> ran.set(true), recorded::set); + + assertTrue("work should have run", ran.get()); + assertTrue("elapsed millis should be recorded (>= 0)", recorded.get() >= 0L); + } + + public void testRecordTimeMillisVoidRecordsTimeEvenOnThrow() { + AtomicLong recorded = new AtomicLong(-1); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows( + IOException.class, + () -> StatsRecorder.recordTimeMillis((StatsRecorder.IORunnable) () -> { throw boom; }, recorded::set) + ); + + assertSame("the original exception must propagate", boom, thrown); + assertTrue("elapsed millis must be recorded even when work throws", recorded.get() >= 0L); + } + + public void testRecordTimeMillisVoidRejectsNullArgs() { + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IORunnable) null, v -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis(() -> {}, null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordTimeMillis(IOSupplier, ...) — value-returning + // ────────────────────────────────────────────────────────────────────── + + public void testRecordTimeMillisValueReturnsResultAndRecordsTime() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + + String result = StatsRecorder.recordTimeMillis(() -> "hello", recorded::set); + + assertEquals("hello", result); + assertTrue("elapsed millis should be recorded", recorded.get() >= 0L); + } + + public void testRecordTimeMillisValuePassesThroughNullResult() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + + String result = StatsRecorder.recordTimeMillis(() -> null, recorded::set); + + assertNull("null result must pass through", result); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordTimeMillisValueRecordsTimeEvenOnThrow() { + AtomicLong recorded = new AtomicLong(-1); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows(IOException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IOSupplier) () -> { + throw boom; + }, recorded::set)); + + assertSame(boom, thrown); + assertTrue("elapsed millis must be recorded even when work throws", recorded.get() >= 0L); + } + + public void testRecordTimeMillisValueRejectsNullArgs() { + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis((StatsRecorder.IOSupplier) null, v -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordTimeMillis(() -> "x", null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordOutcome(IOSupplier, ...) — value-returning + // ────────────────────────────────────────────────────────────────────── + + public void testRecordOutcomeValueSuccessFiresOnSuccessOnly() throws IOException { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + + String result = StatsRecorder.recordOutcome(() -> "ok", recorded::set, success::incrementAndGet, failure::incrementAndGet); + + assertEquals("ok", result); + assertEquals("onSuccess must fire exactly once", 1, success.get()); + assertEquals("onFailure must not fire on success", 0, failure.get()); + assertTrue("time must be recorded", recorded.get() >= 0L); + } + + public void testRecordOutcomeValueFailureFiresOnFailureOnly() { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows(IOException.class, () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) () -> { + throw boom; + }, recorded::set, success::incrementAndGet, failure::incrementAndGet)); + + assertSame("original exception must propagate", boom, thrown); + assertEquals("onFailure must fire exactly once", 1, failure.get()); + assertEquals("onSuccess must not fire on failure", 0, success.get()); + assertTrue("time must be recorded even on failure", recorded.get() >= 0L); + } + + public void testRecordOutcomeValuePropagatesRuntimeException() { + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + RuntimeException boom = new IllegalStateException("rt-boom"); + + RuntimeException thrown = expectThrows( + IllegalStateException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) () -> { + throw boom; + }, v -> {}, success::incrementAndGet, failure::incrementAndGet) + ); + + assertSame(boom, thrown); + assertEquals("onFailure fires on RuntimeException too", 1, failure.get()); + assertEquals(0, success.get()); + } + + public void testRecordOutcomeValueRejectsNullArgs() { + expectThrows( + NullPointerException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IOSupplier) null, v -> {}, () -> {}, () -> {}) + ); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", null, () -> {}, () -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", v -> {}, null, () -> {})); + expectThrows(NullPointerException.class, () -> StatsRecorder.recordOutcome(() -> "x", v -> {}, () -> {}, null)); + } + + // ────────────────────────────────────────────────────────────────────── + // recordOutcome(IORunnable, ...) — void + // ────────────────────────────────────────────────────────────────────── + + public void testRecordOutcomeVoidSuccessFiresOnSuccessOnly() throws IOException { + AtomicBoolean ran = new AtomicBoolean(false); + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + + StatsRecorder.recordOutcome(() -> ran.set(true), recorded::set, success::incrementAndGet, failure::incrementAndGet); + + assertTrue("work should have run", ran.get()); + assertEquals(1, success.get()); + assertEquals(0, failure.get()); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordOutcomeVoidFailureFiresOnFailureOnly() { + AtomicLong recorded = new AtomicLong(-1); + AtomicInteger success = new AtomicInteger(0); + AtomicInteger failure = new AtomicInteger(0); + IOException boom = new IOException("boom"); + + IOException thrown = expectThrows( + IOException.class, + () -> StatsRecorder.recordOutcome( + (StatsRecorder.IORunnable) () -> { throw boom; }, + recorded::set, + success::incrementAndGet, + failure::incrementAndGet + ) + ); + + assertSame(boom, thrown); + assertEquals(1, failure.get()); + assertEquals(0, success.get()); + assertTrue(recorded.get() >= 0L); + } + + public void testRecordOutcomeVoidRejectsNullWork() { + expectThrows( + NullPointerException.class, + () -> StatsRecorder.recordOutcome((StatsRecorder.IORunnable) null, v -> {}, () -> {}, () -> {}) + ); + } + + // ────────────────────────────────────────────────────────────────────── + // Ordering / interaction contracts + // ────────────────────────────────────────────────────────────────────── + + public void testTimeRecordedBeforeOnSuccessCallback() throws IOException { + // Contract: the time recorder runs in finally, before onSuccess is invoked. + AtomicReference order = new AtomicReference<>(""); + + StatsRecorder.recordOutcome( + () -> "v", + millis -> order.updateAndGet(s -> s + "T"), + () -> order.updateAndGet(s -> s + "S"), + () -> order.updateAndGet(s -> s + "F") + ); + + assertEquals("time (T) must be recorded before onSuccess (S)", "TS", order.get()); + } + + public void testOnFailureRunsBeforeExceptionPropagatesAndTimeStillRecorded() { + // Contract: onFailure runs inside catch (before rethrow); time recorder runs in finally. + AtomicReference order = new AtomicReference<>(""); + IOException boom = new IOException("boom"); + + expectThrows( + IOException.class, + () -> StatsRecorder.recordOutcome( + (StatsRecorder.IOSupplier) () -> { throw boom; }, + millis -> order.updateAndGet(s -> s + "T"), + () -> order.updateAndGet(s -> s + "S"), + () -> order.updateAndGet(s -> s + "F") + ) + ); + + // onFailure (F) fires in catch, then time (T) in finally; onSuccess (S) never fires. + assertEquals("expected onFailure then time, no onSuccess", "FT", order.get()); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index 3be2e91b9c609..1bb45186be36f 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -348,32 +348,52 @@ pub struct ShardView { /// Returns a heap-allocated pointer (as i64) to `DataFusionRuntime`. /// Caller must call `close_global_runtime` exactly once to free it. /// -/// # Side effect: spill directory is wiped +/// # Side effect: spill directory contents are wiped (two-phase) /// -/// When `spill_dir` is non-empty, the directory is cleared (`remove_dir_all`) -/// before the `DiskManager` is built. DataFusion does not sweep stale entries -/// itself — `create_local_dirs` in `datafusion-execution` only creates a fresh -/// `datafusion-XXXXXX/` subdirectory and never touches siblings — so a prior -/// non-graceful shutdown (kill -9, OOM-kill, container restart) leaves orphaned -/// `datafusion-*/` trees that accumulate forever. The directory is -/// OpenSearch-owned by contract, so wiping is safe. +/// When `spill_dir` is non-empty, every immediate child is removed: +/// * **Phase 1 (sync, on this thread):** every immediate child (files, +/// symlinks, and subdirectories alike) is renamed to a `.stale` +/// sibling. Renames are metadata-only — one syscall per entry, contents +/// are not touched. +/// * **Phase 2 (async, background thread):** every `*.stale` entry is +/// removed off the boot path. Files and symlinks via `remove_file` +/// (which unlinks the symlink itself without following it); directories +/// via `remove_dir_all`. /// -/// We do not recreate the directory here. DataFusion's `create_local_dirs` will -/// recreate the root via `std::fs::create_dir` and provision a fresh `TempDir` -/// inside it during `DiskManagerBuilder::build()`. +/// The split keeps boot fast even with tens of GB of orphan spill data: the +/// boot thread pays only the rename count (≈10s of µs per entry) and the +/// recursive unlinks happen in parallel with cluster join. Re-scanning by +/// suffix in phase 2 also cleans up `*.stale` leftovers from any prior boot +/// whose cleanup thread did not finish. /// -/// Any cleanup failure (permission denied, read-only filesystem, I/O error, -/// etc.) is operator-actionable and aborts boot with `DataFusionError::Configuration` -/// — the spill directory being in an unrecoverable state implies queries that -/// need to spill would also fail later, so we surface the problem at boot time -/// with full error context instead of letting orphans accumulate silently and -/// failing later mid-query. +/// This sweeps `datafusion-*/` orphans left by a non-graceful shutdown +/// (kill -9, OOM-kill, container restart) which would otherwise accumulate +/// forever — `create_local_dirs` in `datafusion-execution` only creates a +/// fresh `datafusion-XXXXXX/` and never touches siblings. The fresh dir's +/// name has no `.stale` suffix, so phase 2 cannot collide with it. /// -/// Safe today because (a) `datafusion.spill_directory` is `NodeScope + Final`, so -/// `DataFusionPlugin.createComponents` calls this exactly once per JVM in -/// production, and (b) Rust unit tests pass a fresh `tempdir()` per call. Anyone -/// adding a new caller (hot-reload, multiple runtimes sharing a directory, etc.) -/// must rethink this — calling mid-flight will nuke active spill state. +/// The spill directory itself is preserved. It may be a pre-existing mount +/// point whose parent the JVM user does not own; rmdir'ing the root would +/// fail with `EACCES` in that topology. Top-level symlinks are renamed in +/// phase 1 (which renames the link, never the target) and unlinked in +/// phase 2 via `remove_file`, so a stray symlink can never redirect the +/// wipe outside the spill mount. +/// +/// Phase 1 failure aborts boot with `DataFusionError::Configuration` — if +/// we cannot rename orphans now, queries that need to spill would fail +/// later mid-query anyway. Phase 2 failures (including thread-spawn) are +/// logged only; the cluster has either joined or is about to, and +/// stragglers are reaped on the next boot. +/// +/// # Operator contract +/// +/// The OpenSearch process must own `spill_dir` with `rwx`. Write on the +/// parent is **not** required — `spill_dir` is treated as a pre-existing +/// mount point. Callers wanting auto-creation must grant write on the parent. +/// +/// Safe today because `datafusion.spill_directory` is `NodeScope + Final`, so +/// this runs exactly once per JVM. A future caller invoking it mid-flight +/// would nuke active spill state — rethink before adding one. pub fn create_global_runtime( memory_pool_limit: i64, cache_manager_ptr: i64, @@ -410,24 +430,143 @@ pub fn create_global_runtime( spill_limit as u64 }; - // Clear leaked spill files from a prior non-graceful shutdown. Any failure - // here (permission denied, read-only filesystem, I/O error, etc.) is - // operator-actionable and signals that the spill directory is in a state - // where DataFusion cannot reliably write either — fail boot loudly with - // full error context so the operator fixes it before the node joins the - // cluster, instead of accumulating orphan files and inflating - // disk_used_bytes silently. + // Wipe leaked entries from a prior non-graceful shutdown. + // + // Two-phase to keep boot fast even when prior orphans hold tens of GB: + // Phase 1 (sync, on boot thread): rename every immediate child to a + // .stale sibling. One uniform branch — files, symlinks, and + // directories are all renamed identically. Renames are metadata-only + // (one syscall per entry, contents not touched). Failure aborts + // boot loudly with full error context. + // Phase 2 (async, background thread): re-scan and remove every + // *.stale entry — `remove_file` for files and symlinks (so we + // never follow a symlink through to its target), `remove_dir_all` + // for directories. Filtering by the .stale suffix means the fresh + // datafusion-XXXXXX/ that DataFusion provisions next is never a + // deletion target. + // + // The spill directory itself is never removed (mount-point safe). + // Phase 1 only requires write on spill_dir, never on its parent. let spill_path = PathBuf::from(spill_dir); if spill_path.exists() { - if let Err(e) = fs::remove_dir_all(&spill_path) { + let entries = fs::read_dir(&spill_path).map_err(|e| { let msg = format!( - "Failed to clear leaked spill files in {} (io kind={:?}): {}. \ - Verify the OpenSearch process owns the spill directory and the volume \ - is mounted read-write before restarting.", + "Failed to enumerate spill directory {} (io kind={:?}): {}. \ + Verify the process owns the spill directory before restarting.", spill_dir, e.kind(), e ); log::error!("{}", msg); - return Err(DataFusionError::Configuration(msg)); + DataFusionError::Configuration(msg) + })?; + for entry in entries { + let entry = entry.map_err(|e| { + let msg = format!( + "Failed to read spill directory entry in {} (io kind={:?}): {}", + spill_dir, e.kind(), e + ); + log::error!("{}", msg); + DataFusionError::Configuration(msg) + })?; + let path = entry.path(); + let name = entry.file_name(); + let name_str = name.to_string_lossy(); + + // Skip already-renamed leftovers from a prior boot whose async + // cleanup didn't finish — phase 2 below picks them up. + if name_str.ends_with(".stale") { + continue; + } + + // Rename to .stale within the same parent. Metadata-only; + // contents not touched. Same parent (spill_path), so requires + // only write on spill_path itself. Works uniformly for files, + // symlinks, and directories — fs::rename does not follow + // symlinks (it renames the link itself). + let stale = spill_path.join(format!("{}.stale", name_str)); + if let Err(e) = fs::rename(&path, &stale) { + let msg = format!( + "Failed to rename leaked spill entry {} -> {} (io kind={:?}): {}. \ + Verify the process owns the spill directory before restarting.", + path.display(), stale.display(), e.kind(), e + ); + log::error!("{}", msg); + return Err(DataFusionError::Configuration(msg)); + } + } + + // Phase 2: spawn the recursive deletion of all *.stale entries. + // Re-scan inside the thread so we pick up entries from this boot + // AND any *.stale leftovers from prior boots whose cleanup didn't + // finish. Errors inside the thread are logged only — the cluster + // has not yet joined, but stragglers are recoverable next boot. + // + // Spawn failure itself is rare (EAGAIN/ENOMEM under extreme system + // pressure) and not boot-fatal: the disk is already in a valid + // state for DataFusion to start; the *.stale entries just persist + // until a future boot with a healthy thread library reaps them. + // Failing boot here would be more disruptive than the wasted disk. + let spill_path_for_gc = spill_path.clone(); + if let Err(e) = std::thread::Builder::new() + .name("datafusion-spill-gc".to_string()) + .spawn(move || { + let entries = match fs::read_dir(&spill_path_for_gc) { + Ok(e) => e, + Err(e) => { + log::warn!( + "Background spill cleanup: failed to enumerate {} (io kind={:?}): {}", + spill_path_for_gc.display(), e.kind(), e + ); + return; + } + }; + for entry in entries { + let entry = match entry { + Ok(e) => e, + Err(e) => { + log::warn!("Background spill cleanup: read_dir error: {}", e); + continue; + } + }; + let name = entry.file_name(); + if !name.to_string_lossy().ends_with(".stale") { + continue; + } + let path = entry.path(); + // Use file_type (lstat semantics, does not follow symlinks) + // to dispatch: directories use remove_dir_all, everything + // else (regular files, symlinks) uses remove_file. This + // keeps the symlink defense — fs::remove_file on a symlink + // unlinks the link itself without following it. + let result = match entry.file_type() { + Ok(ft) => { + if ft.is_dir() && !ft.is_symlink() { + fs::remove_dir_all(&path) + } else { + fs::remove_file(&path) + } + } + Err(e) => { + log::warn!( + "Background spill cleanup: failed to stat {} (io kind={:?}): {}", + path.display(), e.kind(), e + ); + continue; + } + }; + if let Err(e) = result { + log::warn!( + "Background spill cleanup: failed to remove {} (io kind={:?}): {}", + path.display(), e.kind(), e + ); + } + } + }) + { + log::warn!( + "Failed to spawn background spill cleanup thread: {}. \ + Renamed *.stale entries will accumulate until the next successful boot.", + e + ); } } @@ -810,6 +949,7 @@ pub async unsafe fn fetch_by_row_ids( shard_view.object_metas.as_ref(), shard_view.writer_generations.as_ref(), metadata_cache, + &shard_view.sort_fields, ) .await .map_err(DataFusionError::Execution)?; @@ -1726,6 +1866,19 @@ mod tests { /// tests would race on these globals and produce flaky assertions. static SPILL_GLOBALS_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + /// Test helper: poll until `predicate` returns true or `timeout_ms` elapses. + /// Used to wait on the background spill-cleanup thread without an arbitrary sleep. + fn wait_until bool>(timeout_ms: u64, predicate: F) -> bool { + let deadline = std::time::Instant::now() + std::time::Duration::from_millis(timeout_ms); + while std::time::Instant::now() < deadline { + if predicate() { + return true; + } + std::thread::sleep(std::time::Duration::from_millis(20)); + } + predicate() + } + #[test] fn create_global_runtime_with_empty_spill_dir_disables_disk_manager() { // Empty spill_dir is the "disabled" sentinel from Java. The runtime must build @@ -1773,11 +1926,13 @@ mod tests { let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); assert!(ptr > 0); - // The startup cleanup must have removed the leaked file. - assert!( - !sentinel.exists(), - "leaked spill file must be removed by create_global_runtime startup cleanup" - ); + // Phase 1 renames the sentinel file to leaked_from_prior_run.tmp.stale + // synchronously; phase 2 unlinks it asynchronously. The original name + // is gone immediately; wait for the .stale name to disappear too. + assert!(!sentinel.exists(), "sentinel original name must be gone (renamed)"); + let stale = tmp.path().join("leaked_from_prior_run.tmp.stale"); + let cleaned = wait_until(2000, || !stale.exists()); + assert!(cleaned, "background cleanup must remove the .stale sentinel within 2s"); let runtime = unsafe { &*(ptr as *const DataFusionRuntime) }; assert!( @@ -1797,7 +1952,13 @@ mod tests { // Operator-confirmed contract: the spill directory is OpenSearch-owned and any // contents present at startup are leaked from a prior non-graceful shutdown. // create_global_runtime must clear the directory recursively (files AND - // subdirectories) before constructing the DiskManager. + // subdirectories) before / during constructing the DiskManager. + // + // Cleanup is two-phase: phase 1 renames every immediate child to *.stale on + // the boot thread (uniform across files, symlinks, dirs); phase 2 removes + // them in a background thread (remove_file for files/symlinks, remove_dir_all + // for dirs). The original names are gone immediately after phase 1; wait + // briefly for phase 2 to clear the *.stale entries. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let tmp = tempfile::tempdir().expect("tempdir"); let spill_path = tmp.path().to_str().expect("utf-8 path"); @@ -1810,21 +1971,28 @@ mod tests { fs::create_dir_all(&nested_dir).expect("seed nested subdirs"); let nested_file = nested_dir.join("deep.tmp"); fs::write(&nested_file, b"nested leak").expect("seed nested file"); + let outer_subdir = tmp.path().join("subdir"); assert!(top_file.exists()); assert!(nested_file.exists()); let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); assert!(ptr > 0); - // Both seeded entries must be gone. The directory itself ends up existing again - // because DataFusion's create_local_dirs recreates the root (and provisions a - // fresh datafusion-XXXXXX/ inside it) during DiskManagerBuilder::build() — our - // cleanup only wipes; DataFusion handles recreation. - assert!(!top_file.exists(), "top-level leaked file must be removed"); + // Phase 1: original names gone (renamed to *.stale). + assert!(!top_file.exists(), "top-level file original name must be gone (renamed)"); + assert!(!outer_subdir.exists(), "original subdir name must be gone (renamed)"); + + // Phase 2: *.stale entries cleaned by the background thread. + let stale_top = tmp.path().join("top.tmp.stale"); + let stale_dir = tmp.path().join("subdir.stale"); + let cleaned = wait_until(2000, || !stale_top.exists() && !stale_dir.exists()); + assert!(cleaned, "background cleanup must remove both .stale entries within 2s"); assert!(!nested_file.exists(), "nested leaked file must be removed"); assert!(!nested_dir.exists(), "nested leaked subdir must be removed"); - assert!(tmp.path().exists(), "DataFusion's create_local_dirs must have recreated the root"); - assert!(tmp.path().is_dir(), "spill root must be a directory after DataFusion provisioning"); + + // Spill root itself must remain (cleanup wipes children only). + assert!(tmp.path().exists(), "spill root must be preserved across cleanup"); + assert!(tmp.path().is_dir(), "spill root must remain a directory after cleanup"); unsafe { close_global_runtime(ptr) }; } @@ -1846,8 +2014,8 @@ mod tests { // Cleanup failure is operator-actionable (permissions, RO mount, I/O) and // implies the runtime would fail later mid-query anyway. Surface the failure // at boot with full context. Trigger the failure path by pointing spill_dir - // at a regular file: spill_path.exists() returns true, but remove_dir_all - // refuses to operate on a non-directory. + // at a regular file: spill_path.exists() returns true, but read_dir refuses + // to enumerate a non-directory and returns ErrorKind::NotADirectory. let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let tmp = tempfile::tempdir().expect("tempdir"); let bad_path = tmp.path().join("regular_file"); @@ -1863,12 +2031,202 @@ mod tests { assert!(msg.contains(bad_path_str), "error must reference the offending path; got: {}", msg); assert!(msg.contains("io kind="), "error must include the io::ErrorKind for diagnosis; got: {}", msg); assert!( - msg.contains("Failed to clear leaked spill files"), + msg.contains("Failed to enumerate spill directory") + || msg.contains("Failed to clear leaked spill entry"), "error must identify it as a cleanup failure; got: {}", msg ); } + /// Spill directory is owned by the JVM but its parent isn't (mount-point + /// topology). `fs::remove_dir_all(spill_dir)` would fail at the final + /// rmdir; the fix wipes contents only. + /// + /// Skipped under EUID 0 — root has CAP_DAC_OVERRIDE (or the macOS equivalent) + /// and bypasses the chmod 0o555 we use to simulate the locked parent, so the + /// pre-fix code path would silently succeed and the test would assert nothing. + /// CI runs as a non-root user; if you need to verify under root, run the + /// test in a user namespace or container that drops the capability. + #[test] + #[cfg(unix)] + fn create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent() { + use std::os::unix::fs::PermissionsExt; + + // SAFETY: geteuid() is async-signal-safe and has no preconditions. + let euid = unsafe { libc::geteuid() }; + if euid == 0 { + eprintln!( + "skipping create_global_runtime_succeeds_when_jvm_does_not_own_spill_parent: \ + EUID 0 bypasses chmod 0o555 via DAC override; the pre-fix code path \ + would silently succeed and the test would assert nothing" + ); + return; + } + + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let parent = tempfile::tempdir().expect("parent tempdir"); + let spill_path = parent.path().join("spill"); + fs::create_dir(&spill_path).expect("create spill mount-point dir"); + + // Seed leaked entries: a datafusion-* subtree (kill -9 leftover) and a loose file. + let leaked_dir = spill_path.join("datafusion-aB3kF7"); + fs::create_dir(&leaked_dir).expect("create leaked subdir"); + let leaked_file = leaked_dir.join("tmp_spill_001.arrow"); + fs::write(&leaked_file, b"stale spill data").expect("seed leaked file"); + let loose_file = spill_path.join("stray.txt"); + fs::write(&loose_file, b"loose top-level file").expect("seed loose file"); + + // Lock parent to r+x only — simulates the mount-point parent the JVM doesn't own. + let original_parent_mode = fs::metadata(parent.path()).expect("stat parent").permissions().mode(); + let mut locked = fs::metadata(parent.path()).expect("stat parent").permissions(); + locked.set_mode(0o555); + fs::set_permissions(parent.path(), locked).expect("chmod parent 555"); + + let spill_str = spill_path.to_str().expect("utf-8 path"); + let result = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0); + + // Restore parent perms via RAII so tempdir cleanup runs even on assertion failure. + struct RestorePerms<'a> { + path: &'a std::path::Path, + mode: u32, + } + impl Drop for RestorePerms<'_> { + fn drop(&mut self) { + if let Ok(metadata) = fs::metadata(self.path) { + let mut perms = metadata.permissions(); + perms.set_mode(self.mode); + let _ = fs::set_permissions(self.path, perms); + } + } + } + let _restore = RestorePerms { path: parent.path(), mode: original_parent_mode }; + + let ptr = result.expect("runtime build must succeed when only the parent is read-only"); + assert!(ptr > 0); + + // Phase 1: both originals are renamed inline — gone immediately by the + // original name. Phase 2 unlinks the .stale entries asynchronously. + assert!(!loose_file.exists(), "loose top-level file original name must be gone (renamed)"); + assert!(!leaked_dir.exists(), "leaked datafusion-* original name must be gone (renamed)"); + + let stale_loose = spill_path.join("stray.txt.stale"); + let stale_dir = spill_path.join("datafusion-aB3kF7.stale"); + let cleaned = wait_until(2000, || !stale_loose.exists() && !stale_dir.exists()); + assert!(cleaned, "background cleanup must remove both .stale entries within 2s"); + assert!(!leaked_file.exists(), "leaked file under leaked subdir must be removed"); + + // Spill root itself preserved. + assert!(spill_path.exists(), "spill mount-point dir must be preserved"); + assert!(spill_path.is_dir(), "spill mount-point must remain a directory"); + + unsafe { close_global_runtime(ptr) }; + } + + /// Top-level symlinks must be unlinked, not followed — otherwise the wipe + /// could escape the spill directory and delete files elsewhere. + #[test] + #[cfg(unix)] + fn create_global_runtime_unlinks_top_level_symlink_without_following() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().join("spill"); + fs::create_dir(&spill_path).expect("create spill dir"); + + // outside_target sits next to spill/. If cleanup followed the symlink, it would be deleted. + let outside = tmp.path().join("outside_target.txt"); + fs::write(&outside, b"must NOT be touched").expect("seed outside file"); + + let link = spill_path.join("escape_link"); + std::os::unix::fs::symlink(&outside, &link).expect("create symlink"); + assert!(link.is_symlink(), "precondition: link is a symlink"); + + let spill_str = spill_path.to_str().expect("utf-8 path"); + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_str, 0).expect("runtime build"); + assert!(ptr > 0); + + // Phase 1: symlink is renamed inline (fs::rename does not follow symlinks). + assert!(!link.exists(), "symlink original name must be gone (renamed)"); + let stale_link = spill_path.join("escape_link.stale"); + + // Phase 2: remove_file on a symlink unlinks the link itself, never the target. + let cleaned = wait_until(2000, || !stale_link.exists()); + assert!(cleaned, "background cleanup must remove escape_link.stale within 2s"); + + // Critical: the target outside spill must be intact across both phases. + // If phase 1 had followed the symlink during rename, or phase 2 followed + // it during remove, the outside file would have been clobbered. + assert!( + outside.exists(), + "symlink target outside spill dir must NOT be touched (cleanup must not follow symlinks)" + ); + assert_eq!( + fs::read(&outside).expect("read outside file"), + b"must NOT be touched", + "symlink target contents must be preserved verbatim" + ); + + unsafe { close_global_runtime(ptr) }; + } + + /// Phase 1 (sync rename) must observably move each orphan subdir to a + /// *.stale name. Phase 2 then removes the *.stale entries asynchronously. + /// This pins the rename behavior so a future refactor that goes back to + /// inline recursive removal would be flagged. + #[test] + fn create_global_runtime_renames_orphan_subdirs_to_stale_then_async_removes() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().to_str().expect("utf-8 path"); + + let leaked_a = tmp.path().join("datafusion-aB3kF7"); + fs::create_dir(&leaked_a).expect("create leaked a"); + fs::write(leaked_a.join("tmp_001.arrow"), b"data a").expect("write a"); + let leaked_b = tmp.path().join("datafusion-Xy9pQ2"); + fs::create_dir(&leaked_b).expect("create leaked b"); + fs::write(leaked_b.join("tmp_002.arrow"), b"data b").expect("write b"); + + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + assert!(ptr > 0); + + // Originals were renamed inline — gone immediately by the original name. + assert!(!leaked_a.exists(), "original orphan a must be renamed away"); + assert!(!leaked_b.exists(), "original orphan b must be renamed away"); + + // Phase 2 will eventually remove both *.stale entries. + let stale_a = tmp.path().join("datafusion-aB3kF7.stale"); + let stale_b = tmp.path().join("datafusion-Xy9pQ2.stale"); + let cleaned = wait_until(2000, || !stale_a.exists() && !stale_b.exists()); + assert!(cleaned, "background cleanup must remove both *.stale entries within 2s"); + + unsafe { close_global_runtime(ptr) }; + } + + /// Pre-existing *.stale entries from a prior boot whose phase 2 didn't finish + /// must be cleaned up by the next boot's phase 2, and phase 1 must not + /// double-suffix them (no datafusion-old.stale.stale). + #[test] + fn create_global_runtime_cleans_prior_boot_stale_entries() { + let _guard = SPILL_GLOBALS_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let tmp = tempfile::tempdir().expect("tempdir"); + let spill_path = tmp.path().to_str().expect("utf-8 path"); + + let leftover = tmp.path().join("datafusion-old.stale"); + fs::create_dir(&leftover).expect("create leftover"); + fs::write(leftover.join("residue.arrow"), b"prior boot data").expect("write residue"); + + let ptr = create_global_runtime(64 * 1024 * 1024, 0, spill_path, 0).expect("runtime build"); + assert!(ptr > 0); + + let cleaned = wait_until(2000, || !leftover.exists()); + assert!(cleaned, "prior-boot .stale leftover must be cleaned within 2s"); + assert!( + !tmp.path().join("datafusion-old.stale.stale").exists(), + "phase 1 must NOT double-suffix existing .stale entries" + ); + + unsafe { close_global_runtime(ptr) }; + } + #[test] fn stringview_gc_compacts_sliced_buffers() { let total_rows = 100_000usize; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs index 5ef074d0beff7..d5fb186acbc51 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/cache.rs @@ -6,6 +6,7 @@ * compatible open source license. */ +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use datafusion::execution::cache::cache_manager::{ @@ -28,15 +29,32 @@ fn log_cache_error(operation: &str, error: &str) { // Wrapper to make Mutex implement FileMetadataCache pub struct MutexFileMetadataCache { pub inner: Mutex, + hit_count: AtomicUsize, + miss_count: AtomicUsize, } impl MutexFileMetadataCache { pub fn new(cache: DefaultFilesMetadataCache) -> Self { Self { inner: Mutex::new(cache), + hit_count: AtomicUsize::new(0), + miss_count: AtomicUsize::new(0), } } + pub fn hit_count(&self) -> usize { + self.hit_count.load(Ordering::Relaxed) + } + + pub fn miss_count(&self) -> usize { + self.miss_count.load(Ordering::Relaxed) + } + + pub fn reset_stats(&self) { + self.hit_count.store(0, Ordering::Relaxed); + self.miss_count.store(0, Ordering::Relaxed); + } + pub fn clear_cache(&self) { if let Ok(cache) = self.inner.lock() { cache.clear(); @@ -61,7 +79,15 @@ impl MutexFileMetadataCache { impl CacheAccessor for MutexFileMetadataCache { fn get(&self, k: &Path) -> Option { match self.inner.lock() { - Ok(cache) => cache.get(k), + Ok(cache) => { + let result = cache.get(k); + if result.is_some() { + self.hit_count.fetch_add(1, Ordering::Relaxed); + } else { + self.miss_count.fetch_add(1, Ordering::Relaxed); + } + result + } Err(e) => { log_cache_error("get", &e.to_string()); None diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs index a40bf8f1c0df5..097d3657b9e8e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/custom_cache_manager.rs @@ -504,10 +504,59 @@ impl CustomCacheManager { .unwrap_or(0.0) } + /// Get statistics cache entry count + pub fn statistics_cache_entry_count(&self) -> usize { + self.statistics_cache.as_ref() + .map(|cache| >::len(cache)) + .unwrap_or(0) + } + + /// Get statistics cache size limit in bytes + pub fn statistics_cache_size_limit(&self) -> usize { + self.statistics_cache.as_ref() + .map(|cache| cache.current_size_limit()) + .unwrap_or(0) + } + /// Reset statistics cache stats pub fn statistics_cache_reset_stats(&self) { if let Some(cache) = &self.statistics_cache { cache.reset_stats(); } } + + /// Get metadata cache hit count + pub fn metadata_cache_hit_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.hit_count()) + .unwrap_or(0) + } + + /// Get metadata cache miss count + pub fn metadata_cache_miss_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.miss_count()) + .unwrap_or(0) + } + + /// Get metadata cache entry count + pub fn metadata_cache_entry_count(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| >::len(cache)) + .unwrap_or(0) + } + + /// Get metadata cache size limit in bytes + pub fn metadata_cache_size_limit(&self) -> usize { + self.file_metadata_cache.as_ref() + .map(|cache| cache.get_cache_limit()) + .unwrap_or(0) + } + + /// Reset metadata cache stats + pub fn metadata_cache_reset_stats(&self) { + if let Some(cache) = &self.file_metadata_cache { + cache.reset_stats(); + } + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index edb43e957f60c..2307b0b4356b8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -854,6 +854,7 @@ pub unsafe extern "C" fn df_create_session_context( plan_ptr: *const u8, plan_len: i64, ) -> i64 { + crate::search_stats::inc_listing_table_scan(); let table_name = str_from_raw(table_name_ptr, table_name_len) .map_err(|e| format!("df_create_session_context: {}", e))?; let query_config = @@ -893,6 +894,11 @@ pub unsafe extern "C" fn df_create_session_context_indexed( plan_ptr: *const u8, plan_len: i64, ) -> i64 { + match tree_shape { + 1 => crate::search_stats::inc_single_collector_scan(), + 2 => crate::search_stats::inc_bitmap_tree_scan(), + _ => {} + } let table_name = str_from_raw(table_name_ptr, table_name_len) .map_err(|e| format!("df_create_session_context_indexed: {}", e))?; let query_config = @@ -1156,12 +1162,18 @@ pub unsafe extern "C" fn df_execute_with_context( /// Collects all native executor metrics into a caller-provided byte buffer. /// -/// The buffer must have capacity for at least `size_of::()` bytes (344). +/// `runtime_ptr` may be `0` to skip cache-stats collection. When non-zero it +/// must be a valid pointer returned by [`df_create_global_runtime`]. +/// +/// The buffer must have capacity for at least `size_of::()` bytes (552). /// Returns 0 on success. #[ffm_safe] #[no_mangle] -pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { - use crate::stats::{layout, pack_runtime_metrics, pack_task_monitor, pack_partition_gate, DfStatsBuffer, RuntimeMetricsRepr}; +pub unsafe extern "C" fn df_stats(runtime_ptr: i64, out_ptr: *mut u8, out_cap: i64) -> i64 { + use crate::stats::{ + layout, pack_cache_stats, pack_partition_gate, pack_runtime_metrics, pack_task_monitor, + CacheStatsRepr, DfStatsBuffer, RuntimeMetricsRepr, + }; use crate::task_monitors::{ coordinator_reduce_monitor, query_execution_monitor, stream_next_monitor, plan_setup_monitor, @@ -1190,6 +1202,18 @@ pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { RuntimeMetricsRepr::zeroed() }; + // Cache stats (zeroed when no runtime pointer or no cache manager) + let cache_stats = if runtime_ptr != 0 { + let runtime = &*(runtime_ptr as *const DataFusionRuntime); + runtime + .custom_cache_manager + .as_ref() + .map(pack_cache_stats) + .unwrap_or_else(CacheStatsRepr::default) + } else { + CacheStatsRepr::default() + }; + let buf = DfStatsBuffer { io_runtime, cpu_runtime, @@ -1199,6 +1223,8 @@ pub unsafe extern "C" fn df_stats(out_ptr: *mut u8, out_cap: i64) -> i64 { plan_setup: pack_task_monitor(plan_setup_monitor()), fragment_executor_gate: pack_partition_gate(mgr.cpu_executor.concurrency_gate()), reduce_executor_gate: pack_partition_gate(mgr.coordinator_gate()), + cache_stats, + search_stats: crate::search_stats::snapshot(), }; // Copy struct bytes to caller buffer diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index c22bd03a5b1d3..83c34094d7bdd 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -75,7 +75,7 @@ use crate::api::ShardView; use crate::datafusion_query_config::DatafusionQueryConfig; use crate::indexed_table::bool_tree::residual_bool_to_physical_expr; use crate::indexed_table::metrics::StreamMetrics; -use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruneMetrics}; +use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruneMetrics, StatsPruneTree}; /// Execute an indexed query. @@ -176,6 +176,8 @@ pub async fn execute_indexed_query( table_path: shard_view.table_path.clone(), object_metas: shard_view.object_metas.clone(), writer_generations: shard_view.writer_generations.clone(), + sort_fields: shard_view.sort_fields.clone(), + sort_orders: shard_view.sort_orders.clone(), query_context: crate::query_tracker::QueryTrackingContext::new(context_id, runtime.runtime_env.memory_pool.clone(), crate::query_tracker::QueryType::Shard), table_name: table_name.clone(), indexed_config: None, // derive classification from tree @@ -506,6 +508,8 @@ async unsafe fn execute_indexed_with_context_inner( let table_path = handle.table_path; let object_metas = handle.object_metas; let writer_generations = handle.writer_generations; + let sort_fields = handle.sort_fields; + let sort_orders = handle.sort_orders; let query_context = handle.query_context; let io_handle = handle.io_handle; // Extract context_id early so it can be captured by the per-segment closures @@ -532,6 +536,7 @@ async unsafe fn execute_indexed_with_context_inner( object_metas.as_ref(), writer_generations.as_ref(), metadata_cache, + &sort_fields, ) .await .map_err(DataFusionError::Execution)?; @@ -597,6 +602,22 @@ async unsafe fn execute_indexed_with_context_inner( FilterClass::Tree => None, }; + let prune_tree_config = extraction.as_ref().and_then(|e| { + let mut leaf_exprs: Vec> = Vec::new(); + collect_predicate_exprs(&e.tree, &mut leaf_exprs); + let leaf_predicates: HashMap> = leaf_exprs + .iter() + .filter_map(|expr| { + build_pruning_predicate(expr, schema.clone()) + .map(|pp| (Arc::as_ptr(expr) as *const () as usize, pp)) + }) + .collect(); + if leaf_predicates.is_empty() { + return None; + } + Some((e.tree.clone(), Arc::new(leaf_predicates), schema.clone())) + }); + let predicate_columns = collect_predicate_column_indices(extraction.as_ref()); let factory: EvaluatorFactory = match classification { @@ -615,7 +636,7 @@ async unsafe fn execute_indexed_with_context_inner( .and_then(|expr| build_pruning_predicate(expr, Arc::clone(&schema_for_pruner))); Arc::new( - move |segment: &SegmentFileInfo, _chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, _chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { let pruner = Arc::new(PagePruner::new( &schema_for_pruner, Arc::clone(&segment.metadata), @@ -626,6 +647,7 @@ async unsafe fn execute_indexed_with_context_inner( residual_pruning_predicate.clone(), residual_expr.clone(), Some(PagePruneMetrics::from_stream_metrics(stream_metrics)), + stats_prune_tree.cloned(), )); Ok(eval) }, @@ -691,7 +713,7 @@ async unsafe fn execute_indexed_with_context_inner( let bloom_schema = schema.clone(); let bloom_on_read = query_config.bloom_filter_on_read; Arc::new( - move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { let collector_opt: Option> = match &correctness_provider { Some(provider) => { let collector = FfmSegmentCollector::create( @@ -747,6 +769,7 @@ async unsafe fn execute_indexed_with_context_inner( Arc::new(crate::indexed_table::eval::single_collector::FfmDelegatedBackendCollectorFactory), context_id, bloom_config, + stats_prune_tree.cloned(), )); Ok(eval) }, @@ -799,7 +822,7 @@ async unsafe fn execute_indexed_with_context_inner( ); Arc::new( - move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics| { + move |segment: &SegmentFileInfo, chunk, stream_metrics: &StreamMetrics, stats_prune_tree: Option<&StatsPruneTree>| { // Build one collector per Collector leaf for this chunk. let mut per_leaf: Vec<(i32, Arc)> = Vec::with_capacity(providers.len()); @@ -843,6 +866,7 @@ async unsafe fn execute_indexed_with_context_inner( stream_metrics, )), collector_strategy, + stats_prune_tree: stats_prune_tree.cloned(), }); Ok(eval) }, @@ -870,6 +894,9 @@ async unsafe fn execute_indexed_with_context_inner( query_config: Arc::clone(&query_config), predicate_columns, emit_row_ids, + prune_tree_config, + sort_fields: sort_fields.clone(), + sort_orders: sort_orders.clone(), })); ctx.register_table(&table_name, provider)?; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs index b278dfe706c3a..ce241b62dbcc8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/bitmap_tree.rs @@ -71,7 +71,7 @@ use roaring::RoaringBitmap; use super::{LeafBitmapSource, RgEvalContext, TreeEvaluator, TreePrefetch}; use crate::indexed_table::bool_tree::ResolvedNode; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{packed_bits_to_boolean_array, PositionMap}; use datafusion::physical_optimizer::pruning::PruningPredicate; @@ -88,6 +88,7 @@ impl TreeEvaluator for BitmapTreeEvaluator { page_pruner: &PagePruner, pruning_predicates: &HashMap>, page_prune_metrics: Option<&PagePruneMetrics>, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { let mut per_leaf = Vec::new(); let mut dfs_counter = 0usize; @@ -104,6 +105,7 @@ impl TreeEvaluator for BitmapTreeEvaluator { &mut dfs_counter, &mut per_leaf, /* under_all_and_path */ true, + stats_prune_tree, )?; Ok(TreePrefetch { candidates, @@ -181,7 +183,27 @@ fn prefetch_node( dfs: &mut usize, out: &mut Vec<(usize, RoaringBitmap)>, under_all_and_path: bool, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { + // RG-level subtree pruning: if this subtree provably can't match + // the current RG, skip the entire tree walk. Since collectors are + // always-true in the resolution, a false here means a Predicate + // under AND proved no match — collector bitmaps are irrelevant. + if let Some(spt) = stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(ctx.rg_idx) { + native_bridge_common::log_debug!( + "BitmapTree: skipping subtree for RG {} — pruned by RG-level stats", + ctx.rg_idx + ); + if under_all_and_path { + skip_dfs(node, dfs); + } else { + skip_dfs_with_empty_bitmaps(node, dfs, out); + } + return Ok(RoaringBitmap::new()); + } + } + match node { ResolvedNode::And(children) => { let mut indices: Vec = (0..children.len()).collect(); @@ -207,7 +229,8 @@ fn prefetch_node( page_prune_metrics, dfs, out, - under_all_and_path, // AND preserves the all-AND path + under_all_and_path, + stats_prune_tree.and_then(|spt| spt.children.get(i)), )?; result_bitmap = Some(match result_bitmap { None => child_bitmap, @@ -271,6 +294,7 @@ fn prefetch_node( out, // OR breaks all-AND propagation for its subtree. false, + stats_prune_tree.and_then(|spt| spt.children.get(val)), )?; result_bitmap |= &filtered_bitmap; @@ -304,6 +328,7 @@ fn prefetch_node( dfs, out, /* under_all_and_path */ false, + stats_prune_tree.and_then(|spt| spt.children.first()), )?; // Candidate-stage is a superset. Inverting a superset does // NOT yield a superset of the true NOT — it yields a subset @@ -440,6 +465,36 @@ fn skip_dfs(node: &ResolvedNode, dfs: &mut usize) { } } +/// Advance `dfs` and push empty bitmaps for each Collector leaf without +/// making FFM calls. Used when a subtree is pruned by RG-level stats but +/// refinement may still run (OR/NOT ancestor) — ensures the side-table +/// has entries so refinement doesn't panic on missing keys. +fn skip_dfs_with_empty_bitmaps( + node: &ResolvedNode, + dfs: &mut usize, + out: &mut Vec<(usize, RoaringBitmap)>, +) { + match node { + ResolvedNode::And(children) | ResolvedNode::Or(children) => { + for c in children { + skip_dfs_with_empty_bitmaps(c, dfs, out); + } + } + ResolvedNode::Not(child) => skip_dfs_with_empty_bitmaps(child, dfs, out), + ResolvedNode::Collector { collector, .. } => { + *dfs += 1; + let key = Arc::as_ptr(collector) as *const () as usize; + out.push((key, RoaringBitmap::new())); + } + ResolvedNode::Predicate(_) => *dfs += 1, + ResolvedNode::DelegationPossible { .. } => { + unimplemented!( + "invariant violation: DelegationPossible reached skip_dfs_with_empty_bitmaps." + ) + } + } +} + fn predicate_page_bitmap( expr: &Arc, ctx: &RgEvalContext, @@ -1155,7 +1210,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert_eq!(result.candidates, bm(&[3, 4])); assert_eq!(result.per_leaf.len(), 2); @@ -1169,7 +1224,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert_eq!(result.candidates, bm(&[1, 2, 3])); } @@ -1182,7 +1237,7 @@ mod tests { }; let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // Universe is [0, 16). Minus {0,1,2} = {3..15} let expected: RoaringBitmap = (3u32..16).collect(); @@ -1197,7 +1252,7 @@ mod tests { }; let pruner = empty_pruner(); let state = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -1457,7 +1512,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); assert!(result.candidates.is_empty()); } @@ -1497,7 +1552,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // OR contributes {5} from standalone_leaf → non-empty candidates. assert!(!result.candidates.is_empty()); @@ -1535,7 +1590,7 @@ mod tests { let pruner = empty_pruner(); let result = BitmapTreeEvaluator - .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None) + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) .unwrap(); // NOT inverts empty AND → universe. assert_eq!(result.candidates.len(), 16); @@ -1566,7 +1621,9 @@ mod tests { subtree_cost(&test_predicate_node(), &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE ); - assert_eq!(subtree_cost(&collector_leaf(0), &ctx, &pruner, &pp), ctx.cost_collector * COST_SCALE); + assert_eq!( + subtree_cost(&collector_leaf(0), &ctx, &pruner, &pp), ctx.cost_collector * COST_SCALE + ); } #[test] @@ -1575,7 +1632,9 @@ mod tests { let pruner = empty_pruner(); let pp = HashMap::new(); let wrapped = ResolvedNode::Not(Box::new(test_predicate_node())); - assert_eq!(subtree_cost(&wrapped, &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE); + assert_eq!( + subtree_cost(&wrapped, &ctx, &pruner, &pp), ctx.cost_predicate * COST_SCALE + ); } #[test] @@ -1721,4 +1780,154 @@ mod tests { let bm = RoaringBitmap::new(); assert_eq!(ranges_from_bitmap(&bm, &ctx), vec![]); } + + // ── StatsPruneTree subtree pruning in prefetch ───────────────────── + + fn prune_tree_leaf(can_match: Vec) -> StatsPruneTree { + StatsPruneTree { + rg_can_match: can_match, + children: vec![], + } + } + + fn prune_tree_and( + children: Vec, + ) -> StatsPruneTree { + let mut rg_can_match = vec![true; children[0].rg_can_match.len()]; + for c in &children { + for (r, v) in rg_can_match.iter_mut().zip(c.rg_can_match.iter()) { + *r &= v; + } + } + StatsPruneTree { + rg_can_match, + children, + } + } + + fn prune_tree_or( + children: Vec, + ) -> StatsPruneTree { + let mut rg_can_match = vec![false; children[0].rg_can_match.len()]; + for c in &children { + for (r, v) in rg_can_match.iter_mut().zip(c.rg_can_match.iter()) { + *r |= v; + } + } + StatsPruneTree { + rg_can_match, + children, + } + } + + #[test] + fn stats_prune_tree_root_and_false_skips_entire_rg() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3]), bm(&[3, 4, 5])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_and(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_or_skips_false_child() { + let tree = ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[10, 11, 12]), bm(&[3, 4])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert_eq!(result.candidates, bm(&[3, 4])); + } + + #[test] + fn stats_prune_tree_and_child_false_short_circuits() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3]), bm(&[99])], + }; + let pruner = empty_pruner(); + let spt = prune_tree_and(vec![ + prune_tree_leaf(vec![true]), + prune_tree_leaf(vec![false]), + ]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_nested_or_under_and() { + // AND(OR(collector0, collector1), collector2) + // Cost sort: collector2 (cost=10k) first, OR subtree (cost=20k) second. + // DFS order after sort: collector2=0, OR-child0=1(skipped), OR-child1=2. + let tree = ResolvedNode::And(vec![ + ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]), + collector_leaf(2), + ]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[3, 4, 5]), bm(&[99]), bm(&[3, 4, 5, 6])], + }; + let pruner = empty_pruner(); + let or_spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![true]), + ]); + let spt = prune_tree_and(vec![or_spt, prune_tree_leaf(vec![true])]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + // collector2 (dfs=0) → {3,4,5}; OR-child1 (dfs=2) → {3,4,5,6}; OR = {3,4,5,6} + // AND = {3,4,5} ∩ {3,4,5,6} = {3,4,5} + assert_eq!(result.candidates, bm(&[3, 4, 5])); + } + + #[test] + fn stats_prune_tree_or_all_children_false() { + let tree = ResolvedNode::And(vec![ + ResolvedNode::Or(vec![collector_leaf(0), collector_leaf(1)]), + collector_leaf(2), + ]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2]), bm(&[3, 4]), bm(&[5, 6])], + }; + let pruner = empty_pruner(); + let or_spt = prune_tree_or(vec![ + prune_tree_leaf(vec![false]), + prune_tree_leaf(vec![false]), + ]); + let spt = prune_tree_and(vec![or_spt, prune_tree_leaf(vec![true])]); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, Some(&spt)) + .unwrap(); + assert!(result.candidates.is_empty()); + } + + #[test] + fn stats_prune_tree_none_evaluates_normally() { + let tree = ResolvedNode::And(vec![collector_leaf(0), collector_leaf(1)]); + let leaves = FixedLeafBitmaps { + bitmaps: vec![bm(&[1, 2, 3, 4]), bm(&[3, 4, 5])], + }; + let pruner = empty_pruner(); + let result = BitmapTreeEvaluator + .prefetch(&tree, &test_ctx(), &leaves, &pruner, &HashMap::new(), None, None) + .unwrap(); + assert_eq!(result.candidates, bm(&[3, 4])); + } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs index 49ea62df2cdb2..e00764887b9e1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs @@ -83,6 +83,7 @@ pub(super) fn remap_expr_to_batch( use super::bool_tree::ResolvedNode; use super::page_pruner::PagePruneMetrics; use super::page_pruner::PagePruner; +use super::page_pruner::StatsPruneTree; use super::row_selection::PositionMap; use super::stream::RowGroupInfo; use datafusion::arrow::buffer::Buffer; @@ -272,6 +273,7 @@ pub trait TreeEvaluator: Send + Sync { page_pruner: &PagePruner, pruning_predicates: &HashMap>, page_prune_metrics: Option<&PagePruneMetrics>, + stats_prune_tree: Option<&StatsPruneTree>, ) -> Result; /// Refinement stage: produce the exact per-row `BooleanArray` for one @@ -348,6 +350,8 @@ pub struct TreeBitsetSource { /// recommended here — multiple FFM calls per collector per RG can /// be expensive in multi-collector trees. pub collector_strategy: CollectorCallStrategy, + /// Precomputed per-subtree RG match vectors. Built once at construction. + pub stats_prune_tree: Option, } impl RowGroupBitsetSource for TreeBitsetSource { @@ -358,6 +362,18 @@ impl RowGroupBitsetSource for TreeBitsetSource { max_doc: i32, ) -> Result, String> { let t = Instant::now(); + + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref ann) = self.stats_prune_tree { + if let Some(&false) = ann.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "BitmapTree: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + let ctx = RgEvalContext { rg_idx: rg.index, rg_first_row: rg.first_row, @@ -405,6 +421,7 @@ impl RowGroupBitsetSource for TreeBitsetSource { // inflate counts. We compute final page-level metrics below // after the bitmap tree is fully resolved. None, + self.stats_prune_tree.as_ref(), ) .map_err(|e| format!("TreeBitsetSource::prefetch_rg(rg={}): {}", rg.index, e))?; if prefetch.candidates.is_empty() { @@ -755,6 +772,7 @@ mod tests { _page_pruner: &PagePruner, _pruning_predicates: &HashMap>, _page_prune_metrics: Option<&PagePruneMetrics>, + _stats_prune_tree: Option<&StatsPruneTree>, ) -> Result { Ok(TreePrefetch { candidates: roaring::RoaringBitmap::new(), @@ -803,6 +821,7 @@ mod tests { pruning_predicates: std::sync::Arc::new(HashMap::new()), page_prune_metrics: None, collector_strategy: CollectorCallStrategy::TightenOuterBounds, + stats_prune_tree: None, }; assert!(!source.needs_row_mask()); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs index 8152643673efb..ef76ece7c69b8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs @@ -23,7 +23,7 @@ use roaring::RoaringBitmap; use super::eval_helpers::{compute_page_ranges, evaluate_residual, universe_bitmap_from_page_ranges}; use super::{PrefetchedRg, RowGroupBitsetSource}; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{bitmap_to_packed_bits, PositionMap}; use crate::indexed_table::stream::RowGroupInfo; @@ -35,6 +35,7 @@ pub struct PredicateOnlyEvaluator { pruning_predicate: Option>, residual_expr: Option>, page_prune_metrics: Option, + stats_prune_tree: Option, } impl PredicateOnlyEvaluator { @@ -43,12 +44,14 @@ impl PredicateOnlyEvaluator { pruning_predicate: Option>, residual_expr: Option>, page_prune_metrics: Option, + stats_prune_tree: Option, ) -> Self { Self { page_pruner, pruning_predicate, residual_expr, page_prune_metrics, + stats_prune_tree, } } } @@ -62,6 +65,17 @@ impl RowGroupBitsetSource for PredicateOnlyEvaluator { ) -> Result, String> { let t = Instant::now(); + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref spt) = self.stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "PredicateOnly: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + let page_ranges = compute_page_ranges( self.pruning_predicate.as_ref(), &self.page_pruner, @@ -102,3 +116,66 @@ impl RowGroupBitsetSource for PredicateOnlyEvaluator { Ok(Some(evaluate_residual(residual, batch, batch_len)?)) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::indexed_table::page_pruner::PagePruner; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use datafusion::parquet::arrow::ArrowWriter; + use std::sync::Arc; + use tempfile::NamedTempFile; + + fn minimal_page_pruner() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = datafusion::arrow::record_batch::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![0i32; 8]))], + ) + .unwrap(); + let tmp = NamedTempFile::new().unwrap(); + let mut writer = ArrowWriter::try_new(tmp.reopen().unwrap(), schema.clone(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let file = tmp.reopen().unwrap(); + let options = ArrowReaderOptions::new().with_page_index(true); + let meta = ArrowReaderMetadata::load(&file, options).unwrap(); + Arc::new(PagePruner::new(meta.schema(), meta.metadata().clone())) + } + + #[test] + fn stats_prune_tree_skips_rg_when_false() { + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![false], + children: vec![], + }; + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, Some(spt)); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + assert!(eval.prefetch_rg(&rg, 0, 8).unwrap().is_none()); + } + + #[test] + fn stats_prune_tree_allows_rg_when_true() { + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![true], + children: vec![], + }; + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, Some(spt)); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have candidates"); + assert_eq!(prefetched.candidates.len(), 8); + } + + #[test] + fn stats_prune_tree_none_does_not_prune() { + let pruner = minimal_page_pruner(); + let eval = PredicateOnlyEvaluator::new(pruner, None, None, None, None); + let rg = RowGroupInfo { index: 0, first_row: 0, num_rows: 8 }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have candidates"); + assert_eq!(prefetched.candidates.len(), 8); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs index ab897f4360ff4..5a573dd4db9f6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs @@ -33,7 +33,7 @@ use roaring::RoaringBitmap; use super::{PrefetchedRg, RowGroupBitsetSource}; use crate::indexed_table::ffm_callbacks::{create_provider, FfmSegmentCollector, ProviderHandle}; use crate::indexed_table::index::RowGroupDocsCollector; -use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner, StatsPruneTree}; use crate::indexed_table::row_selection::{ bitmap_to_packed_bits, packed_bits_to_boolean_array, row_selection_to_bitmap, PositionMap, }; @@ -94,7 +94,6 @@ impl DelegatedBackendCollectorFactory for FfmDelegatedBackendCollectorFactory { } } - /// Per-RG state the evaluator keeps for refinement. In row-granular /// mode parquet narrowed fully via `with_predicate` + `RowSelection` /// and nothing is needed here. In block-granular mode we need the @@ -177,6 +176,8 @@ pub struct SingleCollectorEvaluator { context_id: i64, /// Bloom filter pruning config. None = disabled. bloom_config: Option, + /// Precomputed per-RG/subtree match status from RG-level column stats. + stats_prune_tree: Option, } /// Resources needed for per-RG bloom filter pruning. @@ -204,6 +205,7 @@ impl SingleCollectorEvaluator { delegated_backend_collector_factory: Arc, context_id: i64, bloom_config: Option, + stats_prune_tree: Option, ) -> Self { Self { collector, @@ -218,6 +220,7 @@ impl SingleCollectorEvaluator { delegated_backend_collector_factory, context_id, bloom_config, + stats_prune_tree, } } } @@ -258,6 +261,17 @@ impl RowGroupBitsetSource for SingleCollectorEvaluator { ) -> Result, String> { let t = Instant::now(); + // RG-level early-exit: precomputed from column stats at construction. + if let Some(ref spt) = self.stats_prune_tree { + if let Some(&false) = spt.rg_can_match.get(rg.index) { + native_bridge_common::log_debug!( + "SingleCollector: skipping RG {} — pruned by RG-level stats", + rg.index + ); + return Ok(None); + } + } + // Page-prune to discover which row ranges survive. let page_ranges: Option> = self.pruning_predicate.as_ref().and_then(|pp| { self.page_pruner @@ -690,7 +704,7 @@ mod tests { docs: vec![0, 3, 7], }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, @@ -706,7 +720,7 @@ mod tests { fn on_batch_mask_returns_none_for_path_b() { let collector = Arc::new(StubCollector { docs: vec![0] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); let batch = datafusion::arrow::record_batch::RecordBatch::try_new( schema, @@ -734,7 +748,7 @@ mod tests { // (it's the only post-decode filter we have on this path). let collector = Arc::new(StubCollector { docs: vec![0] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); assert!(eval.needs_row_mask()); } @@ -742,7 +756,7 @@ mod tests { fn empty_match_returns_none() { let collector = Arc::new(StubCollector { docs: vec![] }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, first_row: 0, @@ -762,7 +776,7 @@ mod tests { docs: vec![0, 3, 7], }) as Arc; let pruner = minimal_page_pruner(); - let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); let rg = RowGroupInfo { index: 0, @@ -774,8 +788,64 @@ mod tests { assert_eq!(got, vec![0u32, 3, 7]); } + #[test] + fn stats_prune_tree_skips_rg_when_false() { + let collector = Arc::new(StubCollector { + docs: vec![0, 3, 7], + }) as Arc; + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![false], + children: vec![], + }; + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, Some(spt)); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + assert!(eval.prefetch_rg(&rg, 0, 8).unwrap().is_none()); + } + + #[test] + fn stats_prune_tree_allows_rg_when_true() { + let collector = Arc::new(StubCollector { + docs: vec![0, 3, 7], + }) as Arc; + let pruner = minimal_page_pruner(); + let spt = StatsPruneTree { + rg_can_match: vec![true], + children: vec![], + }; + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, Some(spt)); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have matches"); + let got: Vec = prefetched.candidates.iter().collect(); + assert_eq!(got, vec![0u32, 3, 7]); + } + + #[test] + fn stats_prune_tree_none_does_not_prune() { + let collector = Arc::new(StubCollector { + docs: vec![1, 5], + }) as Arc; + let pruner = minimal_page_pruner(); + let eval = SingleCollectorEvaluator::new(Some(collector), pruner, None, None, None, None, CollectorCallStrategy::FullRange, Arc::new(HashMap::new()), 0, Arc::new(FfmDelegatedBackendCollectorFactory), 0, None, None); + let rg = RowGroupInfo { + index: 0, + first_row: 0, + num_rows: 8, + }; + let prefetched = eval.prefetch_rg(&rg, 0, 8).unwrap().expect("should have matches"); + let got: Vec = prefetched.candidates.iter().collect(); + assert_eq!(got, vec![1u32, 5]); + } + // Keep the `fmt` import used #[allow(dead_code)] fn _use(_: &dyn fmt::Debug) {} } - diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs index fc2a6a1254cef..72834911d7057 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/metrics.rs @@ -18,6 +18,8 @@ use datafusion::physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, }; +use crate::indexed_table::parquet_bridge::ReadIoStats; + /// Lightweight metric handles passed from `IndexedExec` to the streaming loop. /// /// All fields are `Option` because standalone uses of `IndexedExec` (i.e. not @@ -100,6 +102,14 @@ pub struct StreamMetrics { /// Time spent polling the inner parquet stream (pull decoded /// batch), isolating decode from our own processing. pub parquet_poll_time: Option